added PersonCollection.get_children()

yamily/__init__.py

@@ -1,3 +1,4 @@
+import collections
 import datetime
 import typing
 
@@ -19,12 +20,19 @@ class Person:
     """
 
     def __init__(self, identifier: str):
-        self.identifier: str = identifier
+        self.__identifier: str = identifier
         self.name: typing.Optional[str] = None
         self.birth_date: typing.Optional[datetime.date] = None
         self.mother: typing.Optional["Person"] = None
         self.father: typing.Optional["Person"] = None
 
+    @property
+    def identifier(self) -> str:
+        return self.__identifier
+
+    def __hash__(self) -> int:
+        return hash(self.__identifier)
+
     def __repr__(self) -> str:
         """
         >>> p = Person("max-mustermann")
@@ -131,6 +139,7 @@ class PersonCollection:
 
     def __init__(self):
         self._persons = {}
+        self.__children = collections.defaultdict(set)
         self.__it = None
 
     def __getitem__(self, identifier: str) -> Person:
@@ -164,8 +173,10 @@ class PersonCollection:
         """
         if person.mother is not None:
             person.mother = self.add_person(person.mother)
+            self.__children[person.mother.identifier].add(person)
         if person.father is not None:
             person.father = self.add_person(person.father)
+            self.__children[person.father.identifier].add(person)
         if person.identifier not in self._persons:
             self._persons[person.identifier] = person
             return person
@@ -173,6 +184,39 @@ class PersonCollection:
         existing_person.merge(person)
         return existing_person
 
+    def get_children(self, parent: typing.Union[Person, str]) -> typing.Set[Person]:
+        """
+        >>> c = PersonCollection()
+
+        >>> alice = Person('alice')
+        >>> alice.father = Person('bob')
+        >>> c.add_person(alice)
+        Person(alice)
+        >>> c.get_children(alice)
+        set()
+        >>> c.get_children('bob')
+        {Person(alice)}
+
+        >>> carol = Person('carol')
+        >>> carol.father = c['bob']
+        >>> c.add_person(carol)
+        Person(carol)
+        >>> c.get_children('bob')
+        {Person(carol), Person(alice)}
+
+        does not support change / removal of parents:
+        >>> carol.father = Person('other-bob')
+        >>> c.add_person(carol)
+        Person(carol)
+        >>> c.get_children('other-bob')
+        {Person(carol)}
+        >>> c.get_children('bob')
+        {Person(carol), Person(alice)}
+        """
+        if isinstance(parent, str):
+            return self.__children[parent]
+        return self.get_children(parent.identifier)
+
     def __iter__(self) -> "PersonCollection":
         """
         >>> c = PersonCollection()