chore: correct returned type

datnguye · Feb 24, 2024 · a03af25 · a03af25
1 parent dc98d01
commit a03af25
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/dbterd/api/__init__.py b/dbterd/api/__init__.py
@@ -1,12 +1,10 @@
 import logging
 from pathlib import Path
-from typing import List, Tuple
 
 from click import Command, Context
 
 from dbterd import default
 from dbterd.adapters.base import Executor
-from dbterd.adapters.meta import Ref, Table
 from dbterd.helpers.log import logger
 
 logger.setLevel(logging.WARNING)  # hide log
@@ -84,7 +82,7 @@ def __set_params_default_if_not_specified(self) -> None:
         self.params["artifacts_dir"] = self.params.get("artifacts_dir", Path.cwd())
         self.params["target"] = self.params.get("target", default.default_target())
 
-    def get_erd(self) -> Tuple[List[Table], List[Ref]]:
+    def get_erd(self) -> str:
         """Generate ERD code for a whole project
 
         Usage:
@@ -94,11 +92,11 @@ def get_erd(self) -> Tuple[List[Table], List[Ref]]:
         ```
 
         Returns:
-            Tuple[List[Table], List[Ref]]: Tables and Refs
+            str: ERD text
         """
         return self.executor.run(**self.params)
 
-    def get_model_erd(self, node_unique_id: str) -> Tuple[List[Table], List[Ref]]:
+    def get_model_erd(self, node_unique_id: str) -> str:
         """Generate ERD code for a model.
 
         Result contains the input model and 1 level relationship model(s) (if any).
@@ -116,6 +114,6 @@ def get_model_erd(self, node_unique_id: str) -> Tuple[List[Table], List[Ref]]:
             - node_unique_id (str): Manifest node unique ID
 
         Returns:
-            Tuple[List[Table], List[Ref]]: Tables and Refs
+            str: ERD text
         """
         return self.executor.run(node_unique_id=node_unique_id, **self.params)