Add more types for function arguments and return values

dbterd/adapters/algos/base.py

@@ -1,4 +1,5 @@
 import copy
+from typing import Dict, List
 
 import click
 
@@ -9,9 +10,10 @@
     TEST_META_RELATIONSHIP_TYPE,
 )
 from dbterd.helpers.log import logger
+from dbterd.types import Catalog, Manifest
 
 
-def get_tables_from_metadata(data=[], **kwargs):
+def get_tables_from_metadata(data=[], **kwargs) -> List[Table]:
     """Extract tables from dbt metadata
 
     Args:
@@ -47,7 +49,7 @@ def get_tables_from_metadata(data=[], **kwargs):
     return tables
 
 
-def get_tables(manifest, catalog, **kwargs):
+def get_tables(manifest: Manifest, catalog: Catalog, **kwargs) -> List[Table]:
     """Extract tables from dbt artifacts
 
     Args:
@@ -94,7 +96,9 @@ def get_tables(manifest, catalog, **kwargs):
     return tables
 
 
-def enrich_tables_from_relationships(tables, relationships):
+def enrich_tables_from_relationships(
+    tables: List[Table], relationships: List[Ref]
+) -> List[Table]:
     """Fullfill columns in Table due to `select *`
 
     Args:
@@ -180,7 +184,7 @@ def get_table_from_metadata(model_metadata, exposures=[], **kwargs) -> Table:
 
 
 def get_table(
-    node_name, manifest_node, catalog_node=None, exposures=[], **kwargs
+    node_name: str, manifest_node, catalog_node=None, exposures=[], **kwargs
 ) -> Table:
     """Construct a single Table object
 
@@ -313,7 +317,7 @@ def get_node_exposures_from_metadata(data=[], **kwargs):
     return exposures
 
 
-def get_node_exposures(manifest):
+def get_node_exposures(manifest: Manifest) -> List[Dict[str, str]]:
     """Get the mapping of table name and exposure name
 
     Args:
@@ -349,7 +353,7 @@ def get_table_name(format: str, **kwargs) -> str:
     return ".".join([kwargs.get(x.lower()) or "KEYNOTFOUND" for x in format.split(".")])
 
 
-def get_relationships_from_metadata(data=[], **kwargs) -> list[Ref]:
+def get_relationships_from_metadata(data=[], **kwargs) -> List[Ref]:
     """Extract relationships from Metadata result list on test relationship
 
     Args:
@@ -409,7 +413,7 @@ def get_relationships_from_metadata(data=[], **kwargs) -> list[Ref]:
     return get_unique_refs(refs=refs)
 
 
-def get_relationships(manifest, **kwargs):
+def get_relationships(manifest: Manifest, **kwargs) -> List[Ref]:
     """Extract relationships from dbt artifacts based on test relationship
 
     Args:
@@ -482,7 +486,7 @@ def get_unique_refs(refs: list[Ref] = []) -> list[Ref]:
     return distinct_list
 
 
-def get_algo_rule(**kwargs):
+def get_algo_rule(**kwargs) -> Dict[str, str]:
     """Extract rule from the --algo option
 
     Args:
@@ -517,7 +521,7 @@ def get_algo_rule(**kwargs):
     return rules
 
 
-def get_table_map_from_metadata(test_node, **kwargs):
+def get_table_map_from_metadata(test_node, **kwargs) -> List[str]:
     """Get the table map with order of [to, from] guaranteed
     (for Metadata)
 
@@ -570,7 +574,7 @@ def get_table_map_from_metadata(test_node, **kwargs):
     return list(reversed(test_parents))
 
 
-def get_table_map(test_node, **kwargs):
+def get_table_map(test_node, **kwargs) -> List[str]:
     """Get the table map with order of [to, from] guaranteed
 
     Args:

dbterd/adapters/algos/test_relationship.py

@@ -1,10 +1,13 @@
+from typing import List, Tuple, Union
+
 from dbterd.adapters.algos import base
 from dbterd.adapters.filter import is_selected_table
-from dbterd.adapters.meta import Ref
+from dbterd.adapters.meta import Ref, Table
 from dbterd.helpers.log import logger
+from dbterd.types import Catalog, Manifest
 
 
-def parse_metadata(data, **kwargs):
+def parse_metadata(manifest: Manifest, **kwargs) -> Tuple[List[Table], List[Ref]]:
     """Get all information (tables, relationships) needed for building diagram
     (from Metadata)
 
@@ -18,7 +21,7 @@ def parse_metadata(data, **kwargs):
     relationships = []
 
     # Parse Table
-    tables = base.get_tables_from_metadata(data=data, **kwargs)
+    tables = base.get_tables_from_metadata(data=manifest, **kwargs)
 
     # Apply selection
     tables = [
@@ -33,7 +36,7 @@ def parse_metadata(data, **kwargs):
     ]
 
     # Parse Ref
-    relationships = base.get_relationships_from_metadata(data=data, **kwargs)
+    relationships = base.get_relationships_from_metadata(data=manifest, **kwargs)
     node_names = [x.node_name for x in tables]
     relationships = [
         x
@@ -47,7 +50,9 @@ def parse_metadata(data, **kwargs):
     return (tables, relationships)
 
 
-def parse(manifest, catalog, **kwargs):
+def parse(
+    manifest: Manifest, catalog: Union[str, Catalog], **kwargs
+) -> Tuple[List[Table], List[Ref]]:
     """Get all information (tables, relationships) needed for building diagram
 
     Args:

dbterd/helpers/file.py

@@ -5,6 +5,7 @@
 from dbt_artifacts_parser import parser
 
 from dbterd.helpers.log import logger
+from dbterd.types import Catalog, Manifest
 
 
 def get_sys_platform():  # pragma: no cover
@@ -109,7 +110,7 @@ def win_prepare_path(path: str) -> str:  # pragma: no cover
     return path
 
 
-def read_manifest(path: str, version: int = None):
+def read_manifest(path: str, version: int = None) -> Manifest:
     """Reads in the manifest.json file, with optional version specification
 
     Args:
@@ -134,7 +135,7 @@ def read_manifest(path: str, version: int = None):
     return parse_func(manifest=_dict)
 
 
-def read_catalog(path: str, version: int = None):
+def read_catalog(path: str, version: int = None) -> Catalog:
     """Reads in the catalog.json file, with optional version specification
 
     Args:

dbterd/types.py

@@ -0,0 +1,31 @@
+from typing import Union
+
+from dbt_artifacts_parser.parsers.catalog.catalog_v1 import CatalogV1
+from dbt_artifacts_parser.parsers.manifest.manifest_v1 import ManifestV1
+from dbt_artifacts_parser.parsers.manifest.manifest_v2 import ManifestV2
+from dbt_artifacts_parser.parsers.manifest.manifest_v3 import ManifestV3
+from dbt_artifacts_parser.parsers.manifest.manifest_v4 import ManifestV4
+from dbt_artifacts_parser.parsers.manifest.manifest_v5 import ManifestV5
+from dbt_artifacts_parser.parsers.manifest.manifest_v6 import ManifestV6
+from dbt_artifacts_parser.parsers.manifest.manifest_v7 import ManifestV7
+from dbt_artifacts_parser.parsers.manifest.manifest_v8 import ManifestV8
+from dbt_artifacts_parser.parsers.manifest.manifest_v9 import ManifestV9
+from dbt_artifacts_parser.parsers.manifest.manifest_v10 import ManifestV10
+from dbt_artifacts_parser.parsers.manifest.manifest_v11 import ManifestV11
+
+Manifest = Union[
+    ManifestV1,
+    ManifestV2,
+    ManifestV3,
+    ManifestV4,
+    ManifestV5,
+    ManifestV6,
+    ManifestV7,
+    ManifestV8,
+    ManifestV9,
+    ManifestV10,
+    ManifestV11,
+]
+
+# If a new version of Catalog is added, replace with `Union[CatalogV1, CatalogV2, ...]`.
+Catalog = CatalogV1