code-climate: mypy: address type-arg

Author: thseiler

strictdoc/export/json/json_generator.py

@@ -1,8 +1,8 @@
-# mypy: disable-error-code="arg-type,attr-defined,no-untyped-def,type-arg,union-attr,operator"
+# mypy: disable-error-code="arg-type,attr-defined,no-untyped-def,union-attr,operator"
 import json
 import os.path
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 from strictdoc.backend.sdoc.models.document import SDocDocument
 from strictdoc.backend.sdoc.models.document_config import DocumentConfig
@@ -17,11 +17,17 @@
 )
 from strictdoc.backend.sdoc.models.section import SDocSection
 from strictdoc.backend.sdoc.models.type_system import (
+    GrammarElementRelationChild,
+    GrammarElementRelationFile,
+    GrammarElementRelationParent,
     RequirementFieldType,
 )
 from strictdoc.core.project_config import ProjectConfig
 from strictdoc.core.traceability_index import TraceabilityIndex
 
+JSONPrimitive = Union[str, int, float, bool, None]
+JSONType = Union[JSONPrimitive, Dict[str, "JSONType"], Sequence["JSONType"]]
+
 
 class TAG(Enum):
     SECTION = 1
@@ -55,7 +61,7 @@ def export_tree(
                 if document_.document_is_included():
                     continue
 
-            document_json_dict = self._write_document(document_)
+            document_json_dict: JSONType = self._write_document(document_)
 
             project_tree_dict["DOCUMENTS"].append(document_json_dict)
 
@@ -65,7 +71,7 @@ def export_tree(
             output_json_file.write(project_tree_json)
 
     @classmethod
-    def _write_document(cls, document: SDocDocument) -> Dict:
+    def _write_document(cls, document: SDocDocument) -> Dict[str, JSONType]:
         document_dict: Dict[str, Any] = {
             "TITLE": document.title,
             "REQ_PREFIX": None,
@@ -160,7 +166,13 @@ def _write_document(cls, document: SDocDocument) -> Dict:
                     cls._write_grammar_field_type(grammar_field)
                 )
 
-            relations: List = element_.relations
+            relations: List[
+                Union[
+                    GrammarElementRelationParent,
+                    GrammarElementRelationChild,
+                    GrammarElementRelationFile,
+                ]
+            ] = element_.relations
             if len(relations) > 0:
                 for element_relation_ in relations:
                     relation_dict = {
@@ -178,7 +190,9 @@ def _write_document(cls, document: SDocDocument) -> Dict:
         return document_dict
 
     @classmethod
-    def _write_node(cls, node, document, level_stack: Optional[Tuple]) -> Dict:
+    def _write_node(
+        cls, node, document, level_stack: Optional[Tuple[int, ...]]
+    ) -> Dict[str, JSONType]:
         def get_level_string_(node_) -> str:
             return (
                 ""
@@ -218,7 +232,7 @@ def get_level_string_(node_) -> str:
             return subnode_dict
 
         elif isinstance(node, SDocDocument):
-            node_dict: Dict[str, List[Dict]] = {JSONKey.NODES: []}
+            node_dict: Dict[str, JSONType] = {JSONKey.NODES: []}
 
             current_number = 0
             for subnode_ in node.section_contents:
@@ -243,9 +257,9 @@ def get_level_string_(node_) -> str:
     @classmethod
     def _write_section(
         cls, section: SDocSection, document: SDocDocument, level_string: str
-    ) -> Dict:
+    ) -> Dict[str, JSONType]:
         assert isinstance(section, (SDocSection, SDocDocument))
-        node_dict: Dict[str, Any] = {
+        node_dict: Dict[str, JSONType] = {
             "_TOC": level_string,
             "TYPE": "SECTION",
             "TITLE": str(section.title),
@@ -272,8 +286,8 @@ def _write_section(
     @classmethod
     def _write_requirement(
         cls, node: SDocNode, document: SDocDocument, level_string: str
-    ) -> Dict:
-        node_dict: Dict[str, Any] = {
+    ) -> Dict[str, JSONType]:
+        node_dict: Dict[str, JSONType] = {
             "_TOC": level_string,
             "TYPE": node.node_type,
         }
@@ -297,7 +311,7 @@ def _write_requirement(
         return node_dict
 
     @classmethod
-    def _write_grammar_field_type(cls, grammar_field) -> Dict:
+    def _write_grammar_field_type(cls, grammar_field) -> Dict[str, JSONType]:
         grammar_field_dict = {
             "TITLE": grammar_field.title,
             "REQUIRED": True if grammar_field.required else False,
@@ -307,8 +321,8 @@ def _write_grammar_field_type(cls, grammar_field) -> Dict:
         return grammar_field_dict
 
     @staticmethod
-    def _write_requirement_relations(node: SDocNode) -> List:
-        relations_list = []
+    def _write_requirement_relations(node: SDocNode) -> Sequence[JSONType]:
+        relations_list: Sequence[JSONType] = []
 
         reference: Reference
         for reference in node.relations: