浏览代码

make person collection iterable

Fabian Peter Hammerle 4 年之前
父节点
当前提交
73944129e4
共有 1 个文件被更改,包括 18 次插入0 次删除
  1. 18 0
      family_tree_yaml/__init__.py

+ 18 - 0
family_tree_yaml/__init__.py

@@ -100,6 +100,7 @@ class Person:
 class PersonCollection:
     def __init__(self):
         self._persons = {}
+        self.__it = None
 
     def __getitem__(self, identifier: str) -> Person:
         """
@@ -132,3 +133,20 @@ class PersonCollection:
         existing_person = self._persons[person.identifier]
         existing_person.merge(person)
         return existing_person
+
+    def __iter__(self) -> "PersonCollection":
+        """
+        >>> c = PersonCollection()
+        >>> for identifier in ['alice', 'bob', 'charlie']:
+        ...     c.add_person(Person(identifier))
+        Person(alice)
+        Person(bob)
+        Person(charlie)
+        >>> list(c)
+        [Person(alice), Person(bob), Person(charlie)]
+        """
+        self.__it = iter(self._persons.values())
+        return self
+
+    def __next__(self) -> Person:
+        return next(self.__it)