ripple-down-rules 0.2.3__py3-none-any.whl → 0.3.0__py3-none-any.whl

ripple_down_rules/prompt.py

@@ -1,97 +1,258 @@
 import ast
 import logging
 import os
+import shutil
+import socket
 import subprocess
 import tempfile
 from _ast import AST
 from functools import cached_property
 from textwrap import indent, dedent
 
-from IPython.core.magic import register_line_magic, line_magic, Magics, magics_class
+from IPython.core.magic import line_magic, Magics, magics_class
 from IPython.terminal.embed import InteractiveShellEmbed
+from colorama import Fore, Style
+from pygments import highlight
+from pygments.formatters.terminal import TerminalFormatter
+from pygments.lexers.python import PythonLexer
 from traitlets.config import Config
-from typing_extensions import List, Optional, Tuple, Dict, Type, Union, Any
+from typing_extensions import List, Optional, Tuple, Dict, Type
 
-from .datastructures.enums import PromptFor
-from .datastructures.case import Case
 from .datastructures.callable_expression import CallableExpression, parse_string_to_expression
+from .datastructures.case import Case
 from .datastructures.dataclasses import CaseQuery
-from .utils import extract_dependencies, contains_return_statement, make_set, get_imports_from_scope, make_list, \
-    get_import_from_type, get_imports_from_types, is_iterable, extract_function_source, encapsulate_user_input, \
-    are_results_subclass_of_types
+from .datastructures.enums import PromptFor, Editor
+from .utils import extract_dependencies, contains_return_statement, get_imports_from_scope, make_list, \
+    get_imports_from_types, extract_function_source, encapsulate_user_input, str_to_snake_case, typing_hint_to_str
+
+
+def detect_available_editor() -> Optional[Editor]:
+    """
+    Detect the available editor on the system.
+
+    :return: The first found editor that is available on the system.
+    """
+    editor_env = os.environ.get("RDR_EDITOR")
+    if editor_env:
+        return Editor.from_str(editor_env)
+    for editor in [Editor.Pycharm, Editor.Code, Editor.CodeServer]:
+        if shutil.which(editor.value):
+            return editor
+    return None
+
+
+def is_port_in_use(port: int = 8080) -> bool:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        return s.connect_ex(("localhost", port)) == 0
+
+
+def start_code_server(workspace):
+    """
+    Start the code-server in the given workspace.
+    """
+    filename = os.path.join(os.path.dirname(__file__), "start-code-server.sh")
+    os.system(f"chmod +x {filename}")
+    print(f"Starting code-server at {filename}")
+    return subprocess.Popen(["/bin/bash", filename, workspace], stdout=subprocess.PIPE,
+                            stderr=subprocess.PIPE, text=True)
 
 
 @magics_class
 class MyMagics(Magics):
-    def __init__(self, shell, scope, output_type: Optional[Type] = None, func_name: str = "user_case",
-                 func_doc: str = "User defined function to be executed on the case.",
-                 code_to_modify: Optional[str] = None):
+    temp_file_path: Optional[str] = None
+    """
+    The path to the temporary file that is created for the user to edit.
+    """
+    port: int = int(os.environ.get("RDR_EDITOR_PORT", 8080))
+    """
+    The port to use for the code-server.
+    """
+    process: Optional[subprocess.Popen] = None
+    """
+    The process of the code-server.
+    """
+
+    def __init__(self, shell, scope,
+                 code_to_modify: Optional[str] = None,
+                 prompt_for: Optional[PromptFor] = None,
+                 case_query: Optional[CaseQuery] = None):
         super().__init__(shell)
         self.scope = scope
-        self.temp_file_path = None
-        self.func_name = func_name
-        self.func_doc = func_doc
         self.code_to_modify = code_to_modify
-        self.output_type = make_list(output_type) if output_type is not None else None
+        self.prompt_for = prompt_for
+        self.case_query = case_query
+        self.output_type = self.get_output_type()
         self.user_edit_line = 0
-        self.function_signature: Optional[str] = None
-        self.build_function_signature()
+        self.func_name: str = self.get_func_name()
+        self.func_doc: str = self.get_func_doc()
+        self.function_signature: str = self.get_function_signature()
+        self.editor: Optional[Editor] = detect_available_editor()
+        self.workspace: str = os.environ.get("RDR_EDITOR_WORKSPACE", os.path.dirname(self.scope['__file__']))
+        self.temp_file_path: str = os.path.join(self.workspace, "edit_code_here.py")
+
+    def get_output_type(self) -> List[Type]:
+        """
+        :return: The output type of the function as a list of types.
+        """
+        if self.prompt_for == PromptFor.Conditions:
+            output_type = bool
+        else:
+            output_type = self.case_query.attribute_type
+        return make_list(output_type) if output_type is not None else None
 
     @line_magic
-    def edit_case(self, line):
+    def edit(self, line):
+        if self.editor is None:
+            print(f"{Fore.RED}ERROR:: No editor found. Please install PyCharm, VSCode or code-server.{Style.RESET_ALL}")
+            return
 
         boilerplate_code = self.build_boilerplate_code()
-
         self.write_to_file(boilerplate_code)
 
-        print(f"Opening {self.temp_file_path} in PyCharm...")
-        subprocess.Popen(["pycharm", "--line", str(self.user_edit_line), self.temp_file_path])
+        self.open_file_in_editor()
+
+    def open_file_in_editor(self):
+        """
+        Open the file in the available editor.
+        """
+        if self.editor == Editor.Pycharm:
+            subprocess.Popen(["pycharm", "--line", str(self.user_edit_line), self.temp_file_path],
+                             stdout=subprocess.DEVNULL,
+                             stderr=subprocess.DEVNULL)
+        elif self.editor == Editor.Code:
+            subprocess.Popen(["code", self.temp_file_path])
+        elif self.editor == Editor.CodeServer:
+            try:
+                subprocess.check_output(["pgrep", "-f", "code-server"])
+                # check if same port is in use
+                if is_port_in_use(self.port):
+                    print(f"Code-server is already running on port {self.port}.")
+                else:
+                    raise ValueError("Port is not in use.")
+            except (subprocess.CalledProcessError, ValueError) as e:
+                self.process = start_code_server(self.workspace)
+            print(f"Open code-server in your browser at http://localhost:{self.port}?folder={self.workspace}")
+        print(f"Edit the file: {Fore.BLUE}{self.temp_file_path}")
 
     def build_boilerplate_code(self):
         imports = self.get_imports()
-        self.build_function_signature()
+        if self.function_signature is None:
+            self.function_signature = self.get_function_signature()
+        if self.func_doc is None:
+            self.func_doc = self.get_func_doc()
         if self.code_to_modify is not None:
             body = indent(dedent(self.code_to_modify), '    ')
         else:
             body = "    # Write your code here\n    pass"
         boilerplate = f"""{imports}\n\n{self.function_signature}\n    \"\"\"{self.func_doc}\"\"\"\n{body}"""
-        self.user_edit_line = imports.count('\n')+6
+        self.user_edit_line = imports.count('\n') + 6
         return boilerplate
 
-    def build_function_signature(self):
-        if self.output_type is None:
-            output_type_hint = ""
-        elif len(self.output_type) == 1:
-            output_type_hint = f" -> {self.output_type[0].__name__}"
+    def get_function_signature(self) -> str:
+        if self.func_name is None:
+            self.func_name = self.get_func_name()
+        output_type_hint = self.get_output_type_hint()
+        func_args = self.get_func_args()
+        return f"def {self.func_name}({func_args}){output_type_hint}:"
+
+    def get_output_type_hint(self) -> str:
+        """
+        :return: A string containing the output type hint for the function.
+        """
+        output_type_hint = ""
+        if self.prompt_for == PromptFor.Conditions:
+            output_type_hint = " -> bool"
+        elif self.prompt_for == PromptFor.Conclusion:
+            output_type_hint = f" -> {self.case_query.attribute_type_hint}"
+        return output_type_hint
+
+    def get_func_args(self) -> str:
+        """
+        :return: A string containing the function arguments.
+        """
+        if self.case_query.is_function:
+            func_args = {}
+            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]
+                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)
+                                   for k, v in func_args.items()])
         else:
-            output_type_hint = f" -> Union[{', '.join([t.__name__ for t in self.output_type])}]"
-        self.function_signature = f"def {self.func_name}(case: {self.case_type.__name__}){output_type_hint}:"
+            func_args = f"case: {self.case_type.__name__}"
+        return func_args
 
     def write_to_file(self, code: str):
-        tmp = tempfile.NamedTemporaryFile(mode='w+', delete=False, suffix=".py",
-                                          dir=os.path.dirname(self.scope['__file__']))
-        tmp.write(code)
-        tmp.flush()
-        self.temp_file_path = tmp.name
-        tmp.close()
+        if self.temp_file_path is None:
+            tmp = tempfile.NamedTemporaryFile(mode='w+', delete=False, suffix=".py",
+                                              dir=self.workspace)
+            tmp.write(code)
+            tmp.flush()
+            self.temp_file_path = tmp.name
+            tmp.close()
+        else:
+            with open(self.temp_file_path, 'w+') as f:
+                f.write(code)
 
     def get_imports(self):
-        case_type_import = f"from {self.case_type.__module__} import {self.case_type.__name__}"
+        """
+        :return: A string containing the imports for the function.
+        """
+        case_type_imports = []
+        if self.case_query.is_function:
+            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]}")
+                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__}")
+        else:
+            case_type_imports.append(f"from {self.case_type.__module__} import {self.case_type.__name__}")
         if self.output_type is None:
             output_type_imports = [f"from typing_extensions import Any"]
         else:
             output_type_imports = get_imports_from_types(self.output_type)
             if len(self.output_type) > 1:
                 output_type_imports.append("from typing_extensions import Union")
-        print(output_type_imports)
+            if list in self.output_type:
+                output_type_imports.append("from typing_extensions import List")
         imports = get_imports_from_scope(self.scope)
         imports = [i for i in imports if ("get_ipython" not in i)]
-        if case_type_import not in imports:
-            imports.append(case_type_import)
+        imports.extend(case_type_imports)
         imports.extend([oti for oti in output_type_imports if oti not in imports])
         imports = set(imports)
         return '\n'.join(imports)
 
+    def get_func_doc(self) -> Optional[str]:
+        """
+        :return: A string containing the function docstring.
+        """
+        if self.prompt_for == PromptFor.Conditions:
+            return (f"Get conditions on whether it's possible to conclude a value"
+                    f" for {self.case_query.name}")
+        else:
+            return f"Get possible value(s) for {self.case_query.name}"
+
+    def get_func_name(self) -> Optional[str]:
+        func_name = ""
+        if self.prompt_for == PromptFor.Conditions:
+            func_name = f"{self.prompt_for.value.lower()}_for_"
+        case_name = self.case_query.name.replace(".", "_")
+        if self.case_query.is_function:
+            # convert any CamelCase word into snake_case by adding _ before each capital letter
+            case_name = case_name.replace(f"_{self.case_query.attribute_name}", "")
+        func_name += case_name
+        return str_to_snake_case(func_name)
+
     @cached_property
     def case_type(self) -> Type:
         """
@@ -103,9 +264,9 @@ class MyMagics(Magics):
         return case._obj_type if isinstance(case, Case) else type(case)
 
     @line_magic
-    def load_case(self, line):
+    def load(self, line):
         if not self.temp_file_path:
-            print("No file to load. Run %edit_case first.")
+            print(f"{Fore.RED}ERROR:: No file to load. Run %edit first.{Style.RESET_ALL}")
             return
 
         with open(self.temp_file_path, 'r') as f:
@@ -118,20 +279,42 @@ class MyMagics(Magics):
                 exec(source, self.scope, exec_globals)
                 user_function = exec_globals[self.func_name]
                 self.shell.user_ns[self.func_name] = user_function
-                print(f"Loaded `{self.func_name}` function into user namespace.")
+                print(f"{Fore.BLUE}Loaded `{self.func_name}` function into user namespace.{Style.RESET_ALL}")
                 return
 
-        print(f"Function `{self.func_name}` not found.")
+        print(f"{Fore.RED}ERROR:: Function `{self.func_name}` not found.{Style.RESET_ALL}")
+
+    @line_magic
+    def help(self, line):
+        """
+        Display help information for the Ipython shell.
+        """
+        help_text = f"""
+Directly write python code in the shell, and then `{Fore.GREEN}return {Fore.RESET}output`. Or use 
+the magic commands to write the code in a temporary file and edit it in PyCharm:
+{Fore.MAGENTA}Usage: %edit{Style.RESET_ALL}
+Opens a temporary file in PyCharm for editing a function (conclusion or conditions for case)
+ that will be executed on the case object.
+{Fore.MAGENTA}Usage: %load{Style.RESET_ALL}
+Loads the function defined in the temporary file into the user namespace, that can then be used inside the
+ Ipython shell. You can then do `{Fore.GREEN}return {Fore.RESET}function_name(case)`.
+        """
+        print(help_text)
+
+    def __del__(self):
+        if hasattr(self, 'process') and self.process is not None and self.process.poll() is None:
+            self.process.terminate()  # Graceful shutdown
+            self.process.wait()  # Ensure cleanup
 
 
 class CustomInteractiveShell(InteractiveShellEmbed):
-    def __init__(self, output_type: Union[Type, Tuple[Type], None] = None, func_name: Optional[str] = None,
-                 func_doc: Optional[str] = None, code_to_modify: Optional[str] = None, **kwargs):
+    def __init__(self, code_to_modify: Optional[str] = None,
+                 prompt_for: Optional[PromptFor] = None,
+                 case_query: Optional[CaseQuery] = None,
+                 **kwargs):
         super().__init__(**kwargs)
-        keys = ['output_type', 'func_name', 'func_doc', 'code_to_modify']
-        values = [output_type, func_name, func_doc, code_to_modify]
-        magics_kwargs = {key: value for key, value in zip(keys, values) if value is not None}
-        self.my_magics = MyMagics(self, self.user_ns, **magics_kwargs)
+        self.my_magics = MyMagics(self, self.user_ns, code_to_modify=code_to_modify,
+                                  prompt_for=prompt_for, case_query=case_query)
         self.register_magics(self.my_magics)
         self.all_lines = []
 
@@ -145,7 +328,6 @@ class CustomInteractiveShell(InteractiveShellEmbed):
                                                          self.my_magics.func_name,
                                                          join_lines=False)[self.my_magics.func_name]
             self.all_lines.append(raw_cell)
-            print("Exiting shell on `return` statement.")
             self.history_manager.store_inputs(line_num=self.execution_count, source=raw_cell)
             self.ask_exit()
             return None
@@ -161,30 +343,23 @@ class IPythonShell:
     """
 
     def __init__(self, scope: Optional[Dict] = None, header: Optional[str] = None,
-                 output_type: Optional[Type] = None, prompt_for: Optional[PromptFor] = None,
-                 attribute_name: Optional[str] = None, attribute_type: Optional[Type] = None,
+                 prompt_for: Optional[PromptFor] = None, case_query: Optional[CaseQuery] = None,
                  code_to_modify: Optional[str] = None):
         """
         Initialize the Ipython shell with the given scope and header.
 
         :param scope: The scope to use for the shell.
         :param header: The header to display when the shell is started.
-        :param output_type: The type of the output from user input.
         :param prompt_for: The type of information to ask the user about.
-        :param attribute_name: The name of the attribute of the case.
-        :param attribute_type: The type of the attribute of the case.
+        :param case_query: The case query which contains the case and the attribute to ask about.
         :param code_to_modify: The code to modify. If given, will be used as a start for user to modify.
         """
         self.scope: Dict = scope or {}
         self.header: str = header or ">>> Embedded Ipython Shell"
-        self.output_type: Optional[Type] = output_type
+        self.case_query: Optional[CaseQuery] = case_query
         self.prompt_for: Optional[PromptFor] = prompt_for
-        self.attribute_name: Optional[str] = attribute_name
-        self.attribute_type: Optional[Type] = attribute_type
         self.code_to_modify: Optional[str] = code_to_modify
         self.user_input: Optional[str] = None
-        self.func_name: str = ""
-        self.func_doc: str = ""
         self.shell: CustomInteractiveShell = self._init_shell()
         self.all_code_lines: List[str] = []
 
@@ -193,49 +368,13 @@ class IPythonShell:
         Initialize the Ipython shell with a custom configuration.
         """
         cfg = Config()
-        self.build_func_name_and_doc()
         shell = CustomInteractiveShell(config=cfg, user_ns=self.scope, banner1=self.header,
-                                       output_type=self.output_type, func_name=self.func_name, func_doc=self.func_doc,
-                                       code_to_modify=self.code_to_modify)
+                                       code_to_modify=self.code_to_modify,
+                                       prompt_for=self.prompt_for,
+                                       case_query=self.case_query,
+                                       )
         return shell
 
-    def build_func_name_and_doc(self) -> Tuple[str, str]:
-        """
-        Build the function name and docstring for the user-defined function.
-
-        :return: A tuple containing the function name and docstring.
-        """
-        case = self.scope['case']
-        case_type = case._obj_type if isinstance(case, Case) else type(case)
-        self.func_name = self.build_func_name(case_type)
-        self.func_doc = self.build_func_doc(case_type)
-
-    def build_func_doc(self, case_type: Type):
-        if self.prompt_for == PromptFor.Conditions:
-            func_doc = (f"Get conditions on whether it's possible to conclude a value"
-                        f" for {case_type.__name__}.{self.attribute_name}")
-        else:
-            func_doc = f"Get possible value(s) for {case_type.__name__}.{self.attribute_name}"
-        if is_iterable(self.attribute_type):
-            possible_types = [t.__name__ for t in self.attribute_type if t not in [list, set]]
-            func_doc += f" of types list/set of {' and/or '.join(possible_types)}"
-        else:
-            func_doc += f" of type {self.attribute_type.__name__}"
-        return func_doc
-
-    def build_func_name(self, case_type: Type):
-        func_name = f"get_{self.prompt_for.value.lower()}_for"
-        func_name += f"_{case_type.__name__}"
-        if self.attribute_name is not None:
-            func_name += f"_{self.attribute_name}"
-        if is_iterable(self.attribute_type):
-            output_names = [f"{t.__name__}" for t in self.attribute_type if t not in [list, set]]
-        else:
-            output_names = [self.attribute_type.__name__] if self.attribute_type is not None else None
-        if output_names is not None:
-            func_name += '_of_type_' + '_'.join(output_names)
-        return func_name.lower()
-
     def run(self):
         """
         Run the embedded shell.
@@ -247,7 +386,7 @@ class IPythonShell:
                 break
             except Exception as e:
                 logging.error(e)
-                print(e)
+                print(f"{Fore.RED}ERROR::{e}{Style.RESET_ALL}")
 
     def update_user_input_from_code_lines(self):
         """
@@ -262,12 +401,16 @@ class IPythonShell:
             else:
                 self.user_input = '\n'.join(self.all_code_lines)
                 self.user_input = encapsulate_user_input(self.user_input, self.shell.my_magics.function_signature,
-                                                         self.func_doc)
-                if f"return {self.func_name}(case)" not in self.user_input:
-                    self.user_input = self.user_input.strip() + f"\nreturn {self.func_name}(case)"
+                                                         self.shell.my_magics.func_doc)
+                if self.case_query.is_function:
+                    args = "**case"
+                else:
+                    args = "case"
+                if f"return {self.shell.my_magics.func_name}({args})" not in self.user_input:
+                    self.user_input = self.user_input.strip() + f"\nreturn {self.shell.my_magics.func_name}({args})"
 
 
-def prompt_user_for_expression(case_query: CaseQuery, prompt_for: PromptFor, prompt_str: Optional[str] = None)\
+def prompt_user_for_expression(case_query: CaseQuery, prompt_for: PromptFor, prompt_str: Optional[str] = None) \
         -> Tuple[Optional[str], Optional[CallableExpression]]:
     """
     Prompt the user for an executable python expression to the given case query.
@@ -282,26 +425,28 @@ def prompt_user_for_expression(case_query: CaseQuery, prompt_for: PromptFor, pro
     while True:
         user_input, expression_tree = prompt_user_about_case(case_query, prompt_for, prompt_str,
                                                              code_to_modify=prev_user_input)
-        prev_user_input = '\n'.join(user_input.split('\n')[2:-1])
         if user_input is None:
             if prompt_for == PromptFor.Conclusion:
-                print("No conclusion provided. Exiting.")
+                print(f"{Fore.YELLOW}No conclusion provided. Exiting.{Style.RESET_ALL}")
                 return None, None
             else:
-                print("Conditions must be provided. Please try again.")
+                print(f"{Fore.RED}Conditions must be provided. Please try again.{Style.RESET_ALL}")
                 continue
+        prev_user_input = '\n'.join(user_input.split('\n')[2:-1])
         conclusion_type = bool if prompt_for == PromptFor.Conditions else case_query.attribute_type
         callable_expression = CallableExpression(user_input, conclusion_type, expression_tree=expression_tree,
-                                                 scope=case_query.scope)
+                                                 scope=case_query.scope,
+                                                  mutually_exclusive=case_query.mutually_exclusive)
         try:
             result = callable_expression(case_query.case)
             if len(make_list(result)) == 0:
-                print(f"The given expression gave an empty result for case {case_query.name}. Please modify!")
+                print(f"{Fore.YELLOW}The given expression gave an empty result for case {case_query.name}."
+                      f" Please modify!{Style.RESET_ALL}")
                 continue
             break
         except Exception as e:
             logging.error(e)
-            print(e)
+            print(f"{Fore.RED}{e}{Style.RESET_ALL}")
     return user_input, callable_expression
 
 
@@ -318,17 +463,26 @@ def prompt_user_about_case(case_query: CaseQuery, prompt_for: PromptFor,
     :return: The user input, and the executable expression that was parsed from the user input.
     """
     if prompt_str is None:
-        prompt_str = f"Give {prompt_for} for {case_query.name}"
+        if prompt_for == PromptFor.Conclusion:
+            prompt_str = f"Give possible value(s) for:\n"
+        else:
+            prompt_str = f"Give conditions on when can the rule be evaluated for:\n"
+        prompt_str += (f"{Fore.CYAN}{case_query.name}{Fore.MAGENTA} of type(s) "
+                       f"{Fore.CYAN}({', '.join(map(lambda x: x.__name__, case_query.core_attribute_type))}){Fore.MAGENTA}")
+        if prompt_for == PromptFor.Conditions:
+            prompt_str += (f"\ne.g. `{Fore.GREEN}return {Fore.BLUE}len{Fore.RESET}(case.attribute) > {Fore.BLUE}0` "
+                           f"{Fore.MAGENTA}\nOR `{Fore.GREEN}return {Fore.YELLOW}True`{Fore.MAGENTA} (If you want the"
+                           f" rule to be always evaluated) \n"
+                           f"You can also do {Fore.YELLOW}%edit{Fore.MAGENTA} for more complex conditions.")
+    prompt_str = f"{Fore.MAGENTA}{prompt_str}{Fore.YELLOW}\n(Write %help for guide){Fore.RESET}"
     scope = {'case': case_query.case, **case_query.scope}
-    output_type = case_query.attribute_type if prompt_for == PromptFor.Conclusion else bool
-    shell = IPythonShell(scope=scope, header=prompt_str, output_type=output_type, prompt_for=prompt_for,
-                         attribute_name=case_query.attribute_name, attribute_type=case_query.attribute_type,
+    shell = IPythonShell(scope=scope, header=prompt_str, prompt_for=prompt_for, case_query=case_query,
                          code_to_modify=code_to_modify)
     return prompt_user_input_and_parse_to_expression(shell=shell)
 
 
 def prompt_user_input_and_parse_to_expression(shell: Optional[IPythonShell] = None,
-                                              user_input: Optional[str] = None)\
+                                              user_input: Optional[str] = None) \
         -> Tuple[Optional[str], Optional[ast.AST]]:
     """
     Prompt the user for input.
@@ -344,11 +498,13 @@ def prompt_user_input_and_parse_to_expression(shell: Optional[IPythonShell] = No
             user_input = shell.user_input
             if user_input is None:
                 return None, None
-            print(user_input)
+            print(f"{Fore.BLUE}Captured User input: {Style.RESET_ALL}")
+            highlighted_code = highlight(user_input, PythonLexer(), TerminalFormatter())
+            print(highlighted_code)
         try:
             return user_input, parse_string_to_expression(user_input)
         except Exception as e:
             msg = f"Error parsing expression: {e}"
             logging.error(msg)
-            print(msg)
+            print(f"{Fore.RED}{msg}{Style.RESET_ALL}")
             user_input = None

ripple_down_rules/rdr.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copyreg
 import importlib
 import sys
 from abc import ABC, abstractmethod
@@ -96,11 +97,12 @@ class RippleDownRules(SubclassJSONSerializer, ABC):
             plt.ioff()
             plt.show()
 
-    def __call__(self, case: Union[Case, SQLTable]) -> CaseAttribute:
+    def __call__(self, case: Union[Case, SQLTable]) -> Union[CallableExpression, Dict[str, CallableExpression]]:
         return self.classify(case)
 
     @abstractmethod
-    def classify(self, case: Union[Case, SQLTable], modify_case: bool = False) -> Optional[CaseAttribute]:
+    def classify(self, case: Union[Case, SQLTable], modify_case: bool = False) \
+            -> Optional[Union[CallableExpression, Dict[str, CallableExpression]]]:
         """
         Classify a case.
 
@@ -111,7 +113,7 @@ class RippleDownRules(SubclassJSONSerializer, ABC):
         pass
 
     def fit_case(self, case_query: CaseQuery, expert: Optional[Expert] = None, **kwargs) \
-            -> Union[CaseAttribute, CallableExpression]:
+            -> Union[CallableExpression, Dict[str, CallableExpression]]:
         """
         Fit the classifier to a case and ask the expert for refinements or alternatives if the classification is
         incorrect by comparing the case with the target category.
@@ -136,7 +138,7 @@ class RippleDownRules(SubclassJSONSerializer, ABC):
 
     @abstractmethod
     def _fit_case(self, case_query: CaseQuery, expert: Optional[Expert] = None, **kwargs) \
-            -> Union[CaseAttribute, CallableExpression]:
+            -> Union[CallableExpression, Dict[str, CallableExpression]]:
         """
         Fit the RDR on a case, and ask the expert for refinements or alternatives if the classification is incorrect by
         comparing the case with the target category.
@@ -929,7 +931,7 @@ class GeneralRDR(RippleDownRules):
         imports += f"from {self.case_type.__module__} import {self.case_type.__name__}\n"
         # add rdr python generated functions.
         for rdr_key, rdr in self.start_rules_dict.items():
-            imports += (f"from {file_path.strip('./')}"
+            imports += (f"from ."
                         f" import {rdr.generated_python_file_name} as {self.rdr_key_to_function_name(rdr_key)}\n")
         return imports