Merge pull request #309 from dkpro/feature/308-Make-type-system-iterable

#308 - Make type system iterable
dkpro · May 1, 2024 · 22232b6 · 22232b6
2 parents fa61c83 + 4495bc4
commit 22232b6
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 0 deletions.
diff --git a/cassis/typesystem.py b/cassis/typesystem.py
@@ -838,6 +838,9 @@ def __init__(self, add_document_annotation_type: bool = True):
         if add_document_annotation_type:
             self._add_document_annotation_type()
 
+    def __iter__(self):
+        return self.get_types()
+
     def contains_type(self, typename: str):
         """Checks whether this type system contains a type with name `typename`.
 

diff --git a/tests/test_typesystem.py b/tests/test_typesystem.py
@@ -511,6 +511,52 @@ def test_is_array():
         assert cas.typesystem.is_array(type.name) == type.name.endswith("Array")
 
 
+def test_get_types():
+    cas = Cas()
+
+    assert {type.name for type in cas.typesystem} == {"uima.tcas.DocumentAnnotation"}
+    assert {type.name for type in cas.typesystem.get_types(built_in=False)} == {"uima.tcas.DocumentAnnotation"}
+    assert {type.name for type in cas.typesystem.get_types(built_in=True)} == {
+        "uima.cas.TOP",
+        "uima.cas.NULL",
+        "uima.cas.Boolean",
+        "uima.cas.Byte",
+        "uima.cas.Short",
+        "uima.cas.Integer",
+        "uima.cas.Long",
+        "uima.cas.Float",
+        "uima.cas.Double",
+        "uima.cas.String",
+        "uima.cas.ArrayBase",
+        "uima.cas.FSArray",
+        "uima.cas.BooleanArray",
+        "uima.cas.ByteArray",
+        "uima.cas.ShortArray",
+        "uima.cas.LongArray",
+        "uima.cas.DoubleArray",
+        "uima.cas.FloatArray",
+        "uima.cas.IntegerArray",
+        "uima.cas.StringArray",
+        "uima.cas.ListBase",
+        "uima.cas.FSList",
+        "uima.cas.EmptyFSList",
+        "uima.cas.NonEmptyFSList",
+        "uima.cas.FloatList",
+        "uima.cas.EmptyFloatList",
+        "uima.cas.NonEmptyFloatList",
+        "uima.cas.IntegerList",
+        "uima.cas.EmptyIntegerList",
+        "uima.cas.NonEmptyIntegerList",
+        "uima.cas.StringList",
+        "uima.cas.EmptyStringList",
+        "uima.cas.NonEmptyStringList",
+        "uima.cas.Sofa",
+        "uima.cas.AnnotationBase",
+        "uima.tcas.Annotation",
+        "uima.tcas.DocumentAnnotation",
+    }
+
+
 @pytest.mark.parametrize(
     "parent_name, child_name, expected",
     [