ripple-down-rules 0.5.5__py3-none-any.whl → 0.5.8__py3-none-any.whl

ripple_down_rules/rdr_decorators.py

@@ -6,16 +6,18 @@ of the RDRs.
 import os.path
 from functools import wraps
 
-from pyparsing.tools.cvt_pyparsing_pep8_names import camel_to_snake
-from typing_extensions import Callable, Optional, Type, Tuple, Dict, Any, Self, get_type_hints, List, Union
+from typing_extensions import Callable, Optional, Type, Tuple, Dict, Any, Self, get_type_hints, List, Union, Sequence
 
-from ripple_down_rules.datastructures.case import create_case, Case
+from ripple_down_rules.datastructures.case import Case
 from ripple_down_rules.datastructures.dataclasses import CaseQuery
-from ripple_down_rules.datastructures.enums import Category
 from ripple_down_rules.experts import Expert, Human
-from ripple_down_rules.rdr import GeneralRDR, RippleDownRules
+from ripple_down_rules.rdr import GeneralRDR
+try:
+    from ripple_down_rules.user_interface.gui import RDRCaseViewer
+except ImportError:
+    RDRCaseViewer = None
 from ripple_down_rules.utils import get_method_args_as_dict, get_func_rdr_model_name, make_set, \
-    get_method_class_if_exists, get_method_name, str_to_snake_case
+    get_method_class_if_exists, str_to_snake_case
 
 
 class RDRDecorator:
@@ -24,15 +26,16 @@ class RDRDecorator:
     def __init__(self, models_dir: str,
                  output_type: Tuple[Type],
                  mutual_exclusive: bool,
-                 python_dir: Optional[str] = None,
                  output_name: str = "output_",
                  fit: bool = True,
-                 expert: Optional[Expert] = None):
+                 expert: Optional[Expert] = None,
+                 ask_always: bool = False,
+                 update_existing_rules: bool = True,
+                 viewer: Optional[RDRCaseViewer] = None):
         """
         :param models_dir: The directory to save/load the RDR models.
         :param output_type: The type of the output. This is used to create the RDR model.
         :param mutual_exclusive: If True, the output types are mutually exclusive.
-        :param python_dir: The directory to save the RDR model as a python file.
          If None, the RDR model will not be saved as a python file.
         :param output_name: The name of the output. This is used to create the RDR model.
         :param fit: If True, the function will be in fit mode. This means that the RDR will prompt the user for the
@@ -40,6 +43,9 @@ class RDRDecorator:
             classification mode. This means that the RDR will classify the function's output based on the RDR model.
         :param expert: The expert that will be used to prompt the user for the correct output. If None, a Human
             expert will be used.
+        :param ask_always: If True, the function will ask the user for a target if it doesn't exist.
+        :param update_existing_rules: If True, the function will update the existing RDR rules
+         even if they gave an output.
         :return: A decorator to use a GeneralRDR as a classifier that monitors and modifies the function's output.
         """
         self.rdr_models_dir = models_dir
@@ -47,10 +53,12 @@ class RDRDecorator:
         self.output_type = output_type
         self.parsed_output_type: List[Type] = []
         self.mutual_exclusive = mutual_exclusive
-        self.rdr_python_path: Optional[str] = python_dir
         self.output_name = output_name
         self.fit: bool = fit
         self.expert: Optional[Expert] = expert
+        self.ask_always = ask_always
+        self.update_existing_rules = update_existing_rules
+        self.viewer = viewer
         self.load()
 
     def decorator(self, func: Callable) -> Callable:
@@ -62,61 +70,77 @@ class RDRDecorator:
                 self.parsed_output_type = self.parse_output_type(func, self.output_type, *args)
             if self.model_name is None:
                 self.initialize_rdr_model_name_and_load(func)
+            if self.expert is None:
+                self.expert = Human(viewer=self.viewer,
+                                    answers_save_path=self.rdr_models_dir + f'/expert_answers')
+
+            func_output = {self.output_name: func(*args, **kwargs)}
 
             if self.fit:
-                expert_answers_path = os.path.join(self.rdr_models_dir, self.model_name, "expert_answers")
-                self.expert = self.expert or Human(answers_save_path=expert_answers_path)
-                case_query = self.create_case_query_from_method(func, self.parsed_output_type,
-                                                                self.mutual_exclusive, self.output_name,
+                case_query = self.create_case_query_from_method(func, func_output,
+                                                                self.parsed_output_type,
+                                                                self.mutual_exclusive,
                                                                 *args, **kwargs)
-                output = self.rdr.fit_case(case_query, expert=self.expert)
+                output = self.rdr.fit_case(case_query, expert=self.expert,
+                                           ask_always_for_target=self.ask_always,
+                                           update_existing_rules=self.update_existing_rules,
+                                           viewer=self.viewer)
+            else:
+                case, case_dict = self.create_case_from_method(func, func_output, *args, **kwargs)
+                output = self.rdr.classify(case)
+
+            if self.output_name in output:
                 return output[self.output_name]
             else:
-                case, case_dict = self.create_case_from_method(func, self.output_name, *args, **kwargs)
-                return self.rdr.classify(case)[self.output_name]
+                return func_output[self.output_name]
 
         return wrapper
 
     @staticmethod
-    def create_case_query_from_method(func: Callable, output_type, mutual_exclusive: bool,
-                                output_name: str = 'output_', *args, **kwargs) -> CaseQuery:
+    def create_case_query_from_method(func: Callable,
+                                      func_output: Dict[str, Any],
+                                      output_type: Sequence[Type],
+                                      mutual_exclusive: bool,
+                                      *args, **kwargs) -> CaseQuery:
         """
         Create a CaseQuery from the function and its arguments.
 
         :param func: The function to create a case from.
-        :param output_type: The type of the output.
+        :param func_output: The output of the function as a dictionary, where the key is the output name.
+        :param output_type: The type of the output as a sequence of types.
         :param mutual_exclusive: If True, the output types are mutually exclusive.
-        :param output_name: The name of the output in the case. Defaults to 'output_'.
         :param args: The positional arguments of the function.
         :param kwargs: The keyword arguments of the function.
         :return: A CaseQuery object representing the case.
         """
         output_type = make_set(output_type)
-        case, case_dict = RDRDecorator.create_case_from_method(func, output_name, *args, **kwargs)
+        case, case_dict = RDRDecorator.create_case_from_method(func, func_output, *args, **kwargs)
         scope = func.__globals__
         scope.update(case_dict)
         func_args_type_hints = get_type_hints(func)
+        output_name = list(func_output.keys())[0]
         func_args_type_hints.update({output_name: Union[tuple(output_type)]})
         return CaseQuery(case, output_name, Union[tuple(output_type)],
                          mutual_exclusive, scope=scope,
                          is_function=True, function_args_type_hints=func_args_type_hints)
 
     @staticmethod
-    def create_case_from_method(func: Callable, output_name: str = "output_", *args, **kwargs) -> Tuple[Case, Dict[str, Any]]:
+    def create_case_from_method(func: Callable,
+                                func_output: Dict[str, Any],
+                                *args, **kwargs) -> Tuple[Case, Dict[str, Any]]:
         """
         Create a Case from the function and its arguments.
 
         :param func: The function to create a case from.
-        :param output_name: The name of the output in the case. Defaults to 'output_'.
+        :param func_output: A dictionary containing the output of the function, where the key is the output name.
         :param args: The positional arguments of the function.
         :param kwargs: The keyword arguments of the function.
         :return: A Case object representing the case.
         """
         case_dict = get_method_args_as_dict(func, *args, **kwargs)
-        func_output = func(*args, **kwargs)
-        case_dict.update({output_name: func_output})
+        case_dict.update(func_output)
         case_name = get_func_rdr_model_name(func)
-        return create_case(case_dict, obj_name=case_name, max_recursion_idx=3), case_dict
+        return Case(dict, id(case_dict), case_name, case_dict, **case_dict), case_dict
 
     def initialize_rdr_model_name_and_load(self, func: Callable) -> None:
         model_file_name = get_func_rdr_model_name(func, include_file_name=True)
@@ -148,10 +172,15 @@ class RDRDecorator:
         """
         Load the RDR model from the specified directory.
         """
-        if self.model_name is not None and os.path.exists(os.path.join(self.rdr_models_dir, self.model_name)):
-            self.rdr = GeneralRDR.load(self.rdr_models_dir, self.model_name)
-        else:
-            self.rdr = GeneralRDR(save_dir=self.rdr_models_dir, model_name=self.model_name)
+        self.rdr = None
+        if self.model_name is not None:
+            model_path = os.path.join(self.rdr_models_dir, self.model_name + f"/rdr_metadata/{self.model_name}.json")
+            if os.path.exists(os.path.join(self.rdr_models_dir, model_path)):
+                self.rdr = GeneralRDR.load(self.rdr_models_dir, self.model_name)
+                self.rdr.set_viewer(self.viewer)
+        if self.rdr is None:
+            self.rdr = GeneralRDR(save_dir=self.rdr_models_dir, model_name=self.model_name,
+                                  viewer=self.viewer)
 
     def update_from_python(self):
         """

ripple_down_rules/rules.py

@@ -3,16 +3,18 @@ from __future__ import annotations
 import logging
 import re
 from abc import ABC, abstractmethod
+from pathlib import Path
 from uuid import uuid4
 
 from anytree import NodeMixin
 from sqlalchemy.orm import DeclarativeBase as SQLTable
-from typing_extensions import List, Optional, Self, Union, Dict, Any, Tuple
+from typing_extensions import List, Optional, Self, Union, Dict, Any, Tuple, Type, Set
 
 from .datastructures.callable_expression import CallableExpression
 from .datastructures.case import Case
+from .datastructures.dataclasses import CaseFactoryMetaData, CaseQuery
 from .datastructures.enums import RDREdge, Stop
-from .utils import SubclassJSONSerializer, conclusion_to_json
+from .utils import SubclassJSONSerializer, conclusion_to_json, get_full_class_name, get_imports_from_types
 
 
 class Rule(NodeMixin, SubclassJSONSerializer, ABC):
@@ -27,7 +29,8 @@ class Rule(NodeMixin, SubclassJSONSerializer, ABC):
                  corner_case: Optional[Union[Case, SQLTable]] = None,
                  weight: Optional[str] = None,
                  conclusion_name: Optional[str] = None,
-                 uid: Optional[str] = None):
+                 uid: Optional[str] = None,
+                 corner_case_metadata: Optional[CaseFactoryMetaData] = None):
         """
         A rule in the ripple down rules classifier.
 
@@ -38,10 +41,13 @@ class Rule(NodeMixin, SubclassJSONSerializer, ABC):
         :param weight: The weight of the rule, which is the type of edge connecting the rule to its parent.
         :param conclusion_name: The name of the conclusion of the rule.
         :param uid: The unique id of the rule.
+        :param corner_case_metadata: Metadata about the corner case, such as the factory that created it or the
+         scenario it is based on.
         """
         super(Rule, self).__init__()
         self.conclusion = conclusion
         self.corner_case = corner_case
+        self.corner_case_metadata: Optional[CaseFactoryMetaData] = corner_case_metadata
         self.parent = parent
         self.weight: Optional[str] = weight
         self.conditions = conditions if conditions else None
@@ -51,6 +57,20 @@ class Rule(NodeMixin, SubclassJSONSerializer, ABC):
         # generate a unique id for the rule using uuid4
         self.uid: str = uid if uid else str(uuid4().int)
 
+    @classmethod
+    def from_case_query(cls, case_query: CaseQuery) -> Rule:
+        """
+        Create a SingleClassRule from a CaseQuery.
+
+        :param case_query: The CaseQuery to create the rule from.
+        :return: A SingleClassRule instance.
+        """
+        corner_case_metadata = CaseFactoryMetaData.from_case_query(case_query)
+        return cls(conditions=case_query.conditions, conclusion=case_query.target,
+                   corner_case=case_query.case, parent=None,
+                   corner_case_metadata=corner_case_metadata,
+                   conclusion_name=case_query.attribute_name)
+
     def _post_detach(self, parent):
         """
         Called after this node is detached from the tree, useful when drawing the tree.
@@ -82,6 +102,33 @@ class Rule(NodeMixin, SubclassJSONSerializer, ABC):
         """
         pass
 
+    def write_corner_case_as_source_code(self, cases_file: str, package_name: Optional[str] = None) -> None:
+        """
+        Write the source code representation of the corner case of the rule to a file.
+
+        :param cases_file: The file to write the corner case to.
+        :param package_name: The package name to use for relative imports.
+        """
+        if self.corner_case_metadata is None:
+            return
+        with open(cases_file, 'a') as f:
+            f.write(f"corner_case_{self.uid} = {self.corner_case_metadata}" + "\n\n\n")
+
+    def get_corner_case_types_to_import(self) -> Set[Type]:
+        """
+        Get the types that need to be imported for the corner case of the rule.
+        """
+        if self.corner_case_metadata is None:
+            return
+        types_to_import = set()
+        if self.corner_case_metadata.factory_method is not None:
+            types_to_import.add(self.corner_case_metadata.factory_method)
+        if self.corner_case_metadata.scenario is not None:
+            types_to_import.add(self.corner_case_metadata.scenario)
+        if self.corner_case_metadata.case_conf is not None:
+            types_to_import.add(self.corner_case_metadata.case_conf)
+        return types_to_import
+
     def write_conclusion_as_source_code(self, parent_indent: str = "", defs_file: Optional[str] = None) -> str:
         """
         Get the source code representation of the conclusion of the rule.
@@ -150,11 +197,16 @@ class Rule(NodeMixin, SubclassJSONSerializer, ABC):
         pass
 
     def _to_json(self) -> Dict[str, Any]:
-        json_serialization = {"conditions": self.conditions.to_json(),
+        try:
+            corner_case = SubclassJSONSerializer.to_json_static(self.corner_case) if self.corner_case else None
+        except Exception as e:
+            logging.debug("Failed to serialize corner case to json, setting it to None. Error: %s", e)
+            corner_case = None
+        json_serialization = {"_type": get_full_class_name(type(self)),
+                              "conditions": self.conditions.to_json(),
                               "conclusion": conclusion_to_json(self.conclusion),
                               "parent": self.parent.json_serialization if self.parent else None,
-                              "corner_case": SubclassJSONSerializer.to_json_static(self.corner_case)
-                              if self.corner_case else None,
+                              "corner_case": corner_case,
                               "conclusion_name": self.conclusion_name,
                               "weight": self.weight,
                               "uid": self.uid}
@@ -277,9 +329,12 @@ class SingleClassRule(Rule, HasAlternativeRule, HasRefinementRule):
             returned_rule = self.alternative(x) if self.alternative else self
         return returned_rule if returned_rule.fired else self
 
-    def fit_rule(self, x: Case, target: CallableExpression, conditions: CallableExpression):
-        new_rule = SingleClassRule(conditions, target,
-                                   corner_case=x, parent=self)
+    def fit_rule(self, case_query: CaseQuery):
+        corner_case_metadata = CaseFactoryMetaData.from_case_query(case_query)
+        new_rule = SingleClassRule(case_query.conditions, case_query.target,
+                                   corner_case=case_query.case, parent=self,
+                                   corner_case_metadata=corner_case_metadata,
+                                   )
         if self.fired:
             self.refinement = new_rule
         else:
@@ -363,11 +418,12 @@ class MultiClassTopRule(Rule, HasRefinementRule, HasAlternativeRule):
         elif self.alternative:  # Here alternative refers to next rule in MultiClassRDR
             return self.alternative
 
-    def fit_rule(self, x: Case, target: CallableExpression, conditions: CallableExpression):
-        if self.fired and target != self.conclusion:
-            self.refinement = MultiClassStopRule(conditions, corner_case=x, parent=self)
+    def fit_rule(self, case_query: CaseQuery):
+        if self.fired and case_query.target != self.conclusion:
+            self.refinement = MultiClassStopRule(case_query.conditions, corner_case=case_query.case, parent=self)
         elif not self.fired:
-            self.alternative = MultiClassTopRule(conditions, target, corner_case=x, parent=self)
+            self.alternative = MultiClassTopRule(case_query.conditions, case_query.target,
+                                                 corner_case=case_query.case, parent=self)
 
     def _to_json(self) -> Dict[str, Any]:
         self.json_serialization = {**Rule._to_json(self),

ripple_down_rules/user_interface/gui.py

@@ -281,12 +281,14 @@ class RDRCaseViewer(QMainWindow):
     main_obj: Optional[Dict[str, Any]] = None
     user_input: Optional[str] = None
     attributes_widget: Optional[QWidget] = None
-    save_function: Optional[Callable[str], None] = None
+    save_function: Optional[Callable[str, str], None] = None
 
-
-    def __init__(self, parent=None, save_file: Optional[str] = None):
+    def __init__(self, parent=None,
+                 save_dir: Optional[str] = None,
+                 save_model_name: Optional[str] = None):
         super().__init__(parent)
-        self.save_file = save_file
+        self.save_dir = save_dir
+        self.save_model_name = save_model_name
 
         self.setWindowTitle("RDR Case Viewer")
 
@@ -323,17 +325,17 @@ class RDRCaseViewer(QMainWindow):
 
         # Add both to main layout
         main_layout.addWidget(self.attributes_widget, stretch=1)
-        main_layout.addWidget(middle_widget, stretch=2)
+        main_layout.addWidget(middle_widget, stretch=1)
         main_layout.addWidget(self.obj_diagram_viewer, stretch=2)
 
-    def set_save_function(self, save_function: Callable[[str], None]) -> None:
+    def set_save_function(self, save_function: Callable[[str, str], None]) -> None:
         """
         Set the function to save the file.
 
         :param save_function: The function to save the file.
         """
         self.save_function = save_function
-        self.save_btn.clicked.connect(lambda: self.save_function(self.save_file))
+        self.save_btn.clicked.connect(lambda: self.save_function(self.save_dir, self.save_model_name))
 
     def print(self, msg):
         """
@@ -489,6 +491,7 @@ class RDRCaseViewer(QMainWindow):
                 self.code_lines, self.template_file_creator.func_name,
                 self.template_file_creator.function_signature,
                 self.template_file_creator.func_doc, self.case_query)
+            self.case_query.scope.update(updates)
         self.template_file_creator = None
 
     def update_attribute_layout(self, obj, name: str):

ripple_down_rules/user_interface/ipython_custom_shell.py

@@ -134,7 +134,7 @@ class IPythonShell:
         """
         Update the user input from the code lines captured in the shell.
         """
-        if len(self.shell.all_lines) == 1 and self.shell.all_lines[0].replace('return', '').strip() == '':
+        if self.shell.all_lines[0].replace('return', '').strip() == '':
             self.user_input = None
         else:
             self.all_code_lines = extract_dependencies(self.shell.all_lines)

ripple_down_rules/user_interface/object_diagram.py

@@ -1,5 +1,9 @@
 import logging
 
+from ripple_down_rules.datastructures.case import Case
+from ripple_down_rules.datastructures.dataclasses import CaseQuery
+from ripple_down_rules.utils import SubclassJSONSerializer
+
 try:
     import graphviz
 except ImportError:
@@ -77,7 +81,11 @@ def generate_object_graph(obj, name='root', seen=None, graph=None, current_depth
         for attr in dir(obj):
             if attr.startswith('_'):
                 continue
-            if attr == 'scope':
+            if isinstance(obj, CaseQuery) and attr == 'scope':
+                continue
+            if isinstance(obj, Case) and attr in ['data']:
+                continue
+            if isinstance(obj, SubclassJSONSerializer) and attr == 'data_class_refs':
                 continue
             value = getattr(obj, attr)
             if callable(value):

ripple_down_rules/user_interface/template_file_creator.py

@@ -8,13 +8,13 @@ from functools import cached_property
 from textwrap import indent, dedent
 
 from colorama import Fore, Style
-from typing_extensions import Optional, Type, List, Callable, Tuple, Dict
+from typing_extensions import Optional, Type, List, Callable, Tuple, Dict, Any, Union
 
 from ..datastructures.case import Case
 from ..datastructures.dataclasses import CaseQuery
 from ..datastructures.enums import Editor, PromptFor
-from ..utils import str_to_snake_case, get_imports_from_scope, make_list, typing_hint_to_str, \
-    get_imports_from_types, extract_function_source, extract_imports
+from ..utils import str_to_snake_case, get_imports_from_scope, make_list, stringify_hint, \
+    get_imports_from_types, extract_function_source, extract_imports, get_types_to_import_from_type_hints
 
 
 def detect_available_editor() -> Optional[Editor]:
@@ -84,6 +84,7 @@ class TemplateFileCreator:
         self.func_doc: str = self.get_func_doc()
         self.function_signature: str = self.get_function_signature()
         self.editor: Optional[Editor] = detect_available_editor()
+        self.editor_cmd: Optional[str] = os.environ.get("RDR_EDITOR_CMD")
         self.workspace: str = os.environ.get("RDR_EDITOR_WORKSPACE", os.path.dirname(self.case_query.scope['__file__']))
         self.temp_file_path: str = os.path.join(self.workspace, "edit_code_here.py")
 
@@ -98,7 +99,7 @@ class TemplateFileCreator:
         return make_list(output_type) if output_type is not None else None
 
     def edit(self):
-        if self.editor is None:
+        if self.editor is None and self.editor_cmd is None:
             self.print_func(
                 f"{Fore.RED}ERROR:: No editor found. Please install PyCharm, VSCode or code-server.{Style.RESET_ALL}")
             return
@@ -112,7 +113,11 @@ class TemplateFileCreator:
         """
         Open the file in the available editor.
         """
-        if self.editor == Editor.Pycharm:
+        if self.editor_cmd is not None:
+            subprocess.Popen([self.editor_cmd, self.temp_file_path],
+                             stdout=subprocess.DEVNULL,
+                             stderr=subprocess.DEVNULL)
+        elif self.editor == Editor.Pycharm:
             subprocess.Popen(["pycharm", "--line", str(self.user_edit_line), self.temp_file_path],
                              stdout=subprocess.DEVNULL,
                              stderr=subprocess.DEVNULL)
@@ -172,7 +177,7 @@ class TemplateFileCreator:
             for k, v in self.case_query.case.items():
                 if (self.case_query.function_args_type_hints is not None
                         and k in self.case_query.function_args_type_hints):
-                    func_args[k] = typing_hint_to_str(self.case_query.function_args_type_hints[k])[0]
+                    func_args[k] = stringify_hint(self.case_query.function_args_type_hints[k])
                 else:
                     func_args[k] = type(v).__name__ if not isinstance(v, type) else f"Type[{v.__name__}]"
             func_args = ', '.join([f"{k}: {v}" if str(v) not in ["NoneType", "None"] else str(k)
@@ -193,7 +198,7 @@ class TemplateFileCreator:
             with open(self.temp_file_path, 'w+') as f:
                 f.write(code)
 
-    def get_imports(self):
+    def get_imports(self) -> str:
         """
         :return: A string containing the imports for the function.
         """
@@ -202,30 +207,25 @@ class TemplateFileCreator:
             for k, v in self.case_query.case.items():
                 if (self.case_query.function_args_type_hints is not None
                         and k in self.case_query.function_args_type_hints):
-                    hint_list = typing_hint_to_str(self.case_query.function_args_type_hints[k])[1]
-                    for hint in hint_list:
-                        hint_split = hint.split('.')
-                        if len(hint_split) > 1:
-                            case_type_imports.append(f"from {'.'.join(hint_split[:-1])} import {hint_split[-1]}")
+                    types_to_import = get_types_to_import_from_type_hints([self.case_query.function_args_type_hints[k]])
+                    case_type_imports.extend(list(types_to_import))
                 else:
-                    if isinstance(v, type):
-                        case_type_imports.append(f"from {v.__module__} import {v.__name__}")
-                    elif hasattr(v, "__module__") and not v.__module__.startswith("__"):
-                        case_type_imports.append(f"\nfrom {type(v).__module__} import {type(v).__name__}")
+                    case_type_imports.append(v)
         else:
-            case_type_imports.append(f"from {self.case_type.__module__} import {self.case_type.__name__}")
+            case_type_imports.append(self.case_type)
         if self.output_type is None:
-            output_type_imports = [f"from typing_extensions import Any"]
+            output_type_imports = [Any]
         else:
-            output_type_imports = get_imports_from_types(self.output_type)
+            output_type_imports = self.output_type
             if len(self.output_type) > 1:
-                output_type_imports.append("from typing_extensions import Union")
+                output_type_imports.append(Union)
             if list in self.output_type:
-                output_type_imports.append("from typing_extensions import List")
-        imports = get_imports_from_scope(self.case_query.scope)
-        imports = [i for i in imports if ("get_ipython" not in i)]
-        imports.extend(case_type_imports)
-        imports.extend([oti for oti in output_type_imports if oti not in imports])
+                output_type_imports.append(List)
+        import_types = list(self.case_query.scope.values())
+        # imports = [i for i in imports if ("get_ipython" not in i)]
+        import_types.extend(case_type_imports)
+        import_types.extend(output_type_imports)
+        imports = get_imports_from_types(import_types)
         imports = set(imports)
         return '\n'.join(imports)