Added additional type checking

marcusklang · Jun 12, 2020 · a799cf7 · a799cf7
1 parent 9cd2b35
commit a799cf7
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/py/docria/model.py b/py/docria/model.py
@@ -1093,21 +1093,26 @@ def bool(default: Optional[bool]=False):
     boolean = bool
 
     @staticmethod
-    def textspan(context="main"):
-        return DataTypeTextspan(DataTypeEnum.SPAN, context=context)
+    def textspan(context: Union[str, Text] = "main"):
+        if isinstance(context, Text):
+            return context.spantype
+        elif isinstance(context, str):
+            return DataTypeTextspan(DataTypeEnum.SPAN, context=context)
+        else:
+            raise ValueError(context)
 
     span = textspan
 
     @staticmethod
-    def noderef(layer):
+    def noderef(layer: str):
         return DataTypeNoderef(DataTypeEnum.NODEREF, layer=layer)
 
     @staticmethod
-    def noderef_many(layer):
+    def noderef_many(layer: str):
         return DataTypeNoderefList(DataTypeEnum.NODEREF_MANY, layer=layer)
 
     @staticmethod
-    def nodespan(layer):
+    def nodespan(layer: str):
         return DataTypeNodespan(DataTypeEnum.NODEREF_SPAN, layer=layer)
 
     @staticmethod