ripple-down-rules 0.6.29__py3-none-any.whl → 0.6.31__py3-none-any.whl

ripple_down_rules/user_interface/gui.py

@@ -1,15 +1,19 @@
-from  __future__ import annotations
+from __future__ import annotations
 
 import inspect
 import logging
+import os.path
 from types import MethodType
 
+from graphviz import Digraph, Source
+
 try:
     from PyQt6.QtCore import Qt
     from PyQt6.QtGui import QPixmap, QPainter, QPalette
     from PyQt6.QtWidgets import (
         QWidget, QVBoxLayout, QLabel, QScrollArea,
-        QSizePolicy, QToolButton, QHBoxLayout, QPushButton, QMainWindow, QGraphicsView, QGraphicsScene, QGraphicsPixmapItem
+        QSizePolicy, QToolButton, QHBoxLayout, QPushButton, QMainWindow, QGraphicsView, QGraphicsScene,
+        QGraphicsPixmapItem
     )
     from qtconsole.inprocess import QtInProcessKernelManager
     from qtconsole.rich_jupyter_widget import RichJupyterWidget
@@ -120,7 +124,7 @@ class BackgroundWidget(QWidget):
 
 
 class CollapsibleBox(QWidget):
-    def __init__(self, title="", parent=None, viewer: Optional[RDRCaseViewer]=None, chain_name: Optional[str] = None,
+    def __init__(self, title="", parent=None, viewer: Optional[RDRCaseViewer] = None, chain_name: Optional[str] = None,
                  main_obj: Optional[Dict[str, Any]] = None):
         super().__init__(parent)
 
@@ -259,12 +263,14 @@ def color_name_to_html(color_name):
         'g': 'green',
         'b': 'blue',
         'o': 'orange',
+        'm': 'magenta',
     }
     color_map = {
         'red': '#d6336c',
         'green': '#2eb872',
         'blue': '#007acc',
         'orange': '#FFA07A',
+        'magenta': '#a22bcf',
     }
     if len(color_name) == 1:
         color_name = single_char_to_name.get(color_name, color_name)
@@ -285,16 +291,10 @@ class RDRCaseViewer(QMainWindow):
     instances: List[RDRCaseViewer] = []
     exit_status: ExitStatus = ExitStatus.CLOSE
 
-    def __init__(self, parent=None,
-                 save_dir: Optional[str] = None,
-                 save_model_name: Optional[str] = None):
+    def __init__(self, parent=None):
         super().__init__(parent)
-        self.exit_status = ExitStatus.CLOSE
         self.instances.clear()
         self.instances.append(self)
-        self.save_dir = save_dir
-        self.save_model_name = save_model_name
-
         self.setWindowTitle("RDR Case Viewer")
 
         self.setBaseSize(1600, 600)  # or your preferred initial size
@@ -333,35 +333,28 @@ class RDRCaseViewer(QMainWindow):
         main_layout.addWidget(middle_widget, stretch=2)
         main_layout.addWidget(self.obj_diagram_viewer, stretch=2)
 
-    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_dir, self.save_model_name))
-
     def print(self, msg):
         """
         Print a message to the console.
         """
         self.ipython_console._append_plain_text(msg + '\n', True)
 
-    def update_for_case_query(self, case_query: CaseQuery, title_txt: Optional[str] = None,
-                              prompt_for: Optional[PromptFor] = None, code_to_modify: Optional[str] = None):
+    def update_for_case_query(self, case_query: CaseQuery, prompt_str: Optional[str] = None,
+                              prompt_for: Optional[PromptFor] = None, code_to_modify: Optional[str] = None,
+                              title: Optional[str] = None):
         self.case_query = case_query
         self.prompt_for = prompt_for
         self.code_to_modify = code_to_modify
-        title_text = title_txt or ""
+        title_text = title or ""
         case_attr_type = ', '.join([t.__name__ for t in case_query.core_attribute_type])
         case_attr_type = style(f"{case_attr_type}", 'g', 28, 'bold')
         case_name = style(f"{case_query.name}", 'b', 28, 'bold')
         title_text = style(f"{title_text} {case_name} of type {case_attr_type}", 'o', 28, 'bold')
-        self.update_for_object(case_query.case, case_query.case_name, case_query.scope, title_text)
+        self.update_for_object(case_query.case, case_query.case_name, scope=case_query.scope, title_text=title_text,
+                               header=prompt_str)
 
     def update_for_object(self, obj: Any, name: str, scope: Optional[dict] = None,
-                          title_text: Optional[str] = None):
+                          title_text: Optional[str] = None, header: Optional[str] = None):
         self.update_main_obj(obj, name)
         title_text = title_text or style(f"{name}", 'o', 28, 'bold')
         scope = scope or {}
@@ -370,6 +363,9 @@ class RDRCaseViewer(QMainWindow):
         self.update_attribute_layout(obj, name)
         self.title_label.setText(title_text)
         self.ipython_console.update_namespace(scope)
+        if header is not None and len(header) > 0:
+            self.ipython_console.print(header)
+        self.exit_status = ExitStatus.CLOSE
 
     def update_main_obj(self, obj, name):
         self.main_obj = {name: obj}
@@ -441,8 +437,6 @@ class RDRCaseViewer(QMainWindow):
                 if isinstance(item.widget(), CollapsibleBox):
                     self.expand_collapse_all(item.widget(), expand=True, curr_depth=curr_depth + 1, max_depth=max_depth)
 
-
-
     def create_buttons_widget(self):
         button_widget = QWidget()
         button_widget_layout = QVBoxLayout(button_widget)
@@ -466,13 +460,19 @@ class RDRCaseViewer(QMainWindow):
         load_btn.clicked.connect(self._load)
         load_btn.setStyleSheet(f"background-color: {color_name_to_html('b')}; color: white;")  # Blue button
 
-        self.save_btn = QPushButton("Save")
-        self.save_btn.setStyleSheet(f"background-color: {color_name_to_html('b')}; color: white;")  # Blue button
+        current_value_btn = QPushButton("Current Value")
+        current_value_btn.clicked.connect(self._show_current_value)
+        current_value_btn.setStyleSheet(f"background-color: {color_name_to_html('m')}; color: white;")
+
+        rule_tree_btn = QPushButton("Rule Tree")
+        rule_tree_btn.clicked.connect(self._show_rule_tree)  # Placeholder for rule tree functionality
+        rule_tree_btn.setStyleSheet(f"background-color: {color_name_to_html('r')}; color: white;")
 
         row_1_button_widget_layout.addWidget(edit_btn)
         row_1_button_widget_layout.addWidget(load_btn)
-        row_1_button_widget_layout.addWidget(accept_btn)
-        row_2_button_widget_layout.addWidget(self.save_btn)
+        row_1_button_widget_layout.addWidget(current_value_btn)
+        row_2_button_widget_layout.addWidget(rule_tree_btn)
+        row_2_button_widget_layout.addWidget(accept_btn)
         return button_widget
 
     def _accept(self):
@@ -481,16 +481,15 @@ class RDRCaseViewer(QMainWindow):
         self.close()
 
     def _edit(self):
-        self.template_file_creator = TemplateFileCreator(self.case_query, self.prompt_for, self.code_to_modify,
-                                                         self.print)
+        self.template_file_creator = self.create_template_file_creator()
         self.template_file_creator.edit()
 
     def _load(self):
         if not self.template_file_creator:
             return
         self.code_lines, updates = self.template_file_creator.load(self.template_file_creator.temp_file_path,
-                                                                    self.template_file_creator.func_name,
-                                                                    self.template_file_creator.print_func)
+                                                                   self.template_file_creator.func_name,
+                                                                   self.template_file_creator.print_func)
         self.ipython_console.kernel.shell.user_ns.update(updates)
         if self.code_lines is not None:
             self.user_input = encapsulate_code_lines_into_a_function(
@@ -498,7 +497,22 @@ class RDRCaseViewer(QMainWindow):
                 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 _show_current_value(self):
+        self.ipython_console.print(self.case_query.current_value_str)
+
+    def _show_rule_tree(self):
+        if self.case_query is None:
+            self.ipython_console.print("No case query provided.")
+            return
+        if not self.case_query.rdr:
+            self.ipython_console.print("No rule tree available for this case query.")
+            return
+        self.case_query.render_rule_tree(view=True)
+
+    def create_template_file_creator(self) -> TemplateFileCreator:
+        return TemplateFileCreator(self.case_query, self.prompt_for, self.code_to_modify,
+                                   self.print)
 
     def update_attribute_layout(self, obj, name: str):
         # Clear the existing layout
@@ -547,7 +561,8 @@ class RDRCaseViewer(QMainWindow):
             attr = f"{attr}"
             try:
                 if is_iterable(value) or hasattr(value, "__dict__") and not inspect.isfunction(value):
-                    self.add_collapsible(attr, value, layout, current_depth + 1, max_depth, chain_name=f"{chain_name}.{attr}")
+                    self.add_collapsible(attr, value, layout, current_depth + 1, max_depth,
+                                         chain_name=f"{chain_name}.{attr}")
                 else:
                     self.add_non_collapsible(attr, value, layout)
             except Exception as e:
@@ -559,7 +574,7 @@ class RDRCaseViewer(QMainWindow):
         type_name = type(value) if not isinstance(value, type) else value
         collapsible = CollapsibleBox(
             f'<b><span style="color:#FFA07A;">{attr}</span></b> {python_colored_repr(type_name)}', viewer=self,
-        chain_name=chain_name, main_obj=self.main_obj)
+            chain_name=chain_name, main_obj=self.main_obj)
         self.add_attributes(value, attr, collapsible.content_layout, current_depth, max_depth, chain_name=chain_name)
         layout.addWidget(collapsible)
 
@@ -587,7 +602,6 @@ class IPythonConsole(RichJupyterWidget):
 
         # Monkey patch its run_cell method
         def custom_run_cell(this, raw_cell, **kwargs):
-            print(raw_cell)
             if contains_return_statement(raw_cell) and 'def ' not in raw_cell:
                 if self.parent.template_file_creator and self.parent.template_file_creator.func_name in raw_cell:
                     self.command_log = self.parent.code_lines
@@ -601,7 +615,6 @@ class IPythonConsole(RichJupyterWidget):
 
         original_run_cell = self.kernel.shell.run_cell
         self.kernel.shell.run_cell = MethodType(custom_run_cell, self.kernel.shell)
-
         self.kernel_client = self.kernel_manager.client()
         self.kernel_client.start_channels()
 
@@ -630,7 +643,7 @@ class IPythonConsole(RichJupyterWidget):
     .in-prompt { font-weight: bold; color: %(in_prompt_color)s }
     .out-prompt-number { font-weight: bold; color: %(out_prompt_number_color)s }
     .out-prompt { font-weight: bold; color: %(out_prompt_color)s }
-'''%dict(
+''' % dict(
             bgcolor='#0b0d0b', fgcolor='#47d9cc', select="#555",
             in_prompt_number_color='lime', in_prompt_color='lime',
             out_prompt_number_color='red', out_prompt_color='red'
@@ -648,6 +661,12 @@ class IPythonConsole(RichJupyterWidget):
         """
         self.kernel.shell.user_ns.update(namespace)
 
+    def print(self, msg):
+        """
+        Custom print function to append messages to the command log.
+        """
+        self.execute(f"print(\"\\n\\n{msg}\")", hidden=True)
+
     def execute(self, source=None, hidden=False, interactive=False):
         # Log the command before execution
         source = source if source is not None else self.input_buffer

ripple_down_rules/user_interface/ipython_custom_shell.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import Optional, List
 
 from IPython.core.magic import magics_class, Magics, line_magic
@@ -6,9 +7,9 @@ from IPython.terminal.embed import InteractiveShellEmbed
 from colorama import Fore, Style
 from traitlets.config import Config
 
+from .template_file_creator import TemplateFileCreator
 from ..datastructures.dataclasses import CaseQuery
 from ..datastructures.enums import PromptFor
-from .template_file_creator import TemplateFileCreator
 from ..utils import contains_return_statement, extract_dependencies, encapsulate_code_lines_into_a_function
 
 
@@ -20,6 +21,7 @@ class MyMagics(Magics):
                  prompt_for: Optional[PromptFor] = None,
                  case_query: Optional[CaseQuery] = None):
         super().__init__(shell)
+        self.case_query: Optional[CaseQuery] = case_query
         self.rule_editor = TemplateFileCreator(case_query, prompt_for=prompt_for, code_to_modify=code_to_modify)
         self.all_code_lines: Optional[List[str]] = None
 
@@ -34,6 +36,29 @@ class MyMagics(Magics):
                                                              self.rule_editor.print_func)
         self.shell.user_ns.update(updates)
 
+    @line_magic
+    def current_value(self, line):
+        """
+        Display the current value of the attribute of the case.
+        """
+        if self.case_query is None:
+            print(f"{Fore.RED}No case query provided.{Style.RESET_ALL}")
+            return
+        print(self.case_query.current_value_str)
+
+    @line_magic
+    def show_rule_tree(self, line):
+        """
+        Display the rule tree for the current case query.
+        """
+        if self.case_query is None:
+            print(f"{Fore.RED}No case query provided.{Style.RESET_ALL}")
+            return
+        if self.case_query.rdr is None:
+            print(f"{Fore.RED}No RDR available for the current case query.{Style.RESET_ALL}")
+            return
+        self.case_query.render_rule_tree(view=True)
+
     @line_magic
     def help(self, line):
         """
@@ -42,12 +67,16 @@ class MyMagics(Magics):
         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}
+{Fore.MAGENTA}%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}
+{Fore.MAGENTA}%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)`.
+{Fore.MAGENTA}%current_value{Style.RESET_ALL}
+Shows the current value of the case attribute on which the rule are being fit.
+{Fore.MAGENTA}%show_rule_tree{Style.RESET_ALL}
+Displays the rule tree for the current case query.
         """
         print(help_text)
 
@@ -144,7 +173,7 @@ class IPythonShell:
                 self.user_input = None
             else:
                 self.user_input = encapsulate_code_lines_into_a_function(self.all_code_lines,
-                                                       function_name=self.shell.my_magics.rule_editor.func_name,
-                                                       function_signature=self.shell.my_magics.rule_editor.function_signature,
-                                                       func_doc=self.shell.my_magics.rule_editor.func_doc,
-                                                       case_query=self.case_query)
+                                                                         function_name=self.shell.my_magics.rule_editor.func_name,
+                                                                         function_signature=self.shell.my_magics.rule_editor.function_signature,
+                                                                         func_doc=self.shell.my_magics.rule_editor.func_doc,
+                                                                         case_query=self.case_query)

ripple_down_rules/user_interface/prompt.py

@@ -1,7 +1,10 @@
 import ast
 import logging
+import sys
 from _ast import AST
 
+from .. import logger
+
 try:
     from PyQt6.QtWidgets import QApplication
     from .gui import RDRCaseViewer, style
@@ -29,7 +32,7 @@ class UserPrompt:
     """
     shell_lock: RLock = RLock()  # To ensure that only one thread can access the shell at a time
 
-    def __init__(self):
+    def __init__(self, prompt_user: bool = True):
         """
         Initialize the UserPrompt class.
         """
@@ -51,18 +54,19 @@ class UserPrompt:
         callable_expression: Optional[CallableExpression] = None
         while True:
             with self.shell_lock:
-                user_input, expression_tree = self.prompt_user_about_case(case_query, prompt_for, prompt_str, code_to_modify=prev_user_input)
+                user_input, expression_tree = self.prompt_user_about_case(case_query, prompt_for, prompt_str,
+                                                                          code_to_modify=prev_user_input)
             if user_input is None:
                 if prompt_for == PromptFor.Conclusion:
-                    self.print_func(f"{Fore.YELLOW}No conclusion provided. Exiting.{Style.RESET_ALL}")
+                    self.print_func(f"\n{Fore.YELLOW}No conclusion provided. Exiting.{Style.RESET_ALL}")
                     return None, None
                 else:
-                    self.print_func(f"{Fore.RED}Conditions must be provided. Please try again.{Style.RESET_ALL}")
+                    self.print_func(f"\n{Fore.RED}Conditions must be provided. Please try again.{Style.RESET_ALL}")
                     continue
-            elif user_input == "exit":
-                self.print_func(f"{Fore.YELLOW}Exiting.{Style.RESET_ALL}")
+            elif user_input in ["exit", 'quit']:
+                self.print_func(f"\n{Fore.YELLOW}Exiting.{Style.RESET_ALL}")
                 return user_input, None
-                
+
             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,
@@ -73,8 +77,9 @@ class UserPrompt:
                 if len(make_list(result)) == 0 and (user_input_to_modify is not None
                                                     and (prev_user_input != user_input_to_modify)):
                     user_input_to_modify = prev_user_input
-                    self.print_func(f"{Fore.YELLOW}The given expression gave an empty result for case {case_query.name}."
-                          f" Please accept or modify!{Style.RESET_ALL}")
+                    self.print_func(
+                        f"{Fore.YELLOW}The given expression gave an empty result for case {case_query.name}."
+                        f" Please accept or modify!{Style.RESET_ALL}")
                     continue
                 break
             except Exception as e:
@@ -82,7 +87,6 @@ class UserPrompt:
                 self.print_func(f"{Fore.RED}{e}{Style.RESET_ALL}")
         return user_input, callable_expression
 
-
     def prompt_user_about_case(self, case_query: CaseQuery, prompt_for: PromptFor,
                                prompt_str: Optional[str] = None,
                                code_to_modify: Optional[str] = None) -> Tuple[Optional[str], Optional[AST]]:
@@ -95,7 +99,7 @@ class UserPrompt:
         :param code_to_modify: The code to modify. If given will be used as a start for user to modify.
         :return: The user input, and the executable expression that was parsed from the user input.
         """
-        self.print_func("Entered shell")
+        logger.debug("Entered shell")
         initial_prompt_str = f"{prompt_str}\n" if prompt_str is not None else ''
         if prompt_for == PromptFor.Conclusion:
             prompt_for_str = f"Give possible value(s) for:"
@@ -103,20 +107,32 @@ class UserPrompt:
             prompt_for_str = f"Give conditions for:"
         case_query.scope.update({'case': case_query.case})
         shell = None
+
         if self.viewer is None:
+            prompt_for_str = prompt_for_str.replace(":", f" {case_query.name}:")
             prompt_str = f"{Fore.WHITE}{initial_prompt_str}{Fore.MAGENTA}{prompt_for_str}"
             prompt_str = self.construct_prompt_str_for_shell(case_query, prompt_for, prompt_str)
             shell = IPythonShell(header=prompt_str, prompt_for=prompt_for, case_query=case_query,
-                                    code_to_modify=code_to_modify)
+                                 code_to_modify=code_to_modify)
         else:
-            prompt_str = initial_prompt_str + prompt_for_str
-            self.viewer.update_for_case_query(case_query, prompt_str,
-                                                prompt_for=prompt_for, code_to_modify=code_to_modify)
+            prompt_str = case_query.current_value_str
+            self.viewer.update_for_case_query(case_query, prompt_for=prompt_for, code_to_modify=code_to_modify,
+                                              title=prompt_for_str, prompt_str=prompt_str)
         user_input, expression_tree = self.prompt_user_input_and_parse_to_expression(shell=shell)
-        self.print_func("Exited shell")
+        logger.debug("Exited shell")
         return user_input, expression_tree
 
-
+    def build_prompt_str_for_ai(self, case_query: CaseQuery, prompt_for: PromptFor,
+                                initial_prompt_str: Optional[str] = None) -> str:
+        initial_prompt_str = f"{initial_prompt_str}\n" if initial_prompt_str is not None else ''
+        if prompt_for == PromptFor.Conclusion:
+            prompt_for_str = f"Give possible value(s) for:"
+        else:
+            prompt_for_str = f"Give conditions for:"
+        prompt_for_str = prompt_for_str.replace(":", f" {case_query.name}:")
+        prompt_str = f"{Fore.WHITE}{initial_prompt_str}{Fore.MAGENTA}{prompt_for_str}"
+        prompt_str += '\n' + case_query.current_value_str
+        return prompt_str
 
     def construct_prompt_str_for_shell(self, case_query: CaseQuery, prompt_for: PromptFor,
                                        prompt_str: Optional[str] = None) -> str:
@@ -127,8 +143,7 @@ class UserPrompt:
         :param prompt_for: The type of information the user should provide for the given case.
         :param prompt_str: The prompt string to display to the user.
         """
-        prompt_str += (f"\n{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}")
+        prompt_str += '\n' + case_query.current_value_str
         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"
@@ -138,7 +153,6 @@ class UserPrompt:
         prompt_str = f"{Fore.MAGENTA}{prompt_str}{Fore.YELLOW}\n(Write %help for guide){Fore.RESET}\n"
         return prompt_str
 
-
     def prompt_user_input_and_parse_to_expression(self, shell: Optional[IPythonShell] = None,
                                                   user_input: Optional[str] = None) \
             -> Tuple[Optional[str], Optional[ast.AST]]:
@@ -152,17 +166,18 @@ class UserPrompt:
         while True:
             if user_input is None:
                 user_input = self.start_shell_and_get_user_input(shell=shell)
-                if user_input is None or user_input == 'exit':
+                if user_input is None or user_input in ['exit', 'quit']:
                     return user_input, None
-                self.print_func(f"{Fore.GREEN}Captured User input: {Style.RESET_ALL}")
-                highlighted_code = highlight(user_input, PythonLexer(), TerminalFormatter())
-                self.print_func(highlighted_code)
+                if logger.level <= logging.DEBUG:
+                    self.print_func(f"\n{Fore.GREEN}Captured User input: {Style.RESET_ALL}")
+                    highlighted_code = highlight(user_input, PythonLexer(), TerminalFormatter())
+                    self.print_func(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)
-                self.print_func(f"{Fore.RED}{msg}{Style.RESET_ALL}")
+                self.print_func(f"\n{Fore.RED}{msg}{Style.RESET_ALL}")
                 user_input = None
 
     def start_shell_and_get_user_input(self, shell: Optional[IPythonShell] = None) -> Optional[str]:
@@ -185,6 +200,6 @@ class UserPrompt:
             self.viewer.show()
             app.exec()
             if self.viewer.exit_status == ExitStatus.CLOSE:
-                exit(0)
+                sys.exit()
             user_input = self.viewer.user_input
         return user_input

ripple_down_rules/user_interface/template_file_creator.py

@@ -109,20 +109,21 @@ class TemplateFileCreator:
 
         self.open_file_in_editor()
 
-    def open_file_in_editor(self):
+    def open_file_in_editor(self, file_path: Optional[str] = None):
         """
         Open the file in the available editor.
         """
+        file_path = file_path or self.temp_file_path
         if self.editor_cmd is not None:
-            subprocess.Popen([self.editor_cmd, self.temp_file_path],
+            subprocess.Popen([self.editor_cmd, 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],
+            subprocess.Popen(["pycharm", "--line", str(self.user_edit_line), file_path],
                              stdout=subprocess.DEVNULL,
                              stderr=subprocess.DEVNULL)
         elif self.editor == Editor.Code:
-            subprocess.Popen(["code", self.temp_file_path])
+            subprocess.Popen(["code", file_path])
         elif self.editor == Editor.CodeServer:
             try:
                 subprocess.check_output(["pgrep", "-f", "code-server"])
@@ -134,7 +135,8 @@ class TemplateFileCreator:
             except (subprocess.CalledProcessError, ValueError) as e:
                 self.process = start_code_server(self.workspace)
             self.print_func(f"Open code-server in your browser at http://localhost:{self.port}?folder={self.workspace}")
-        self.print_func(f"Edit the file: {Fore.MAGENTA}{self.temp_file_path}")
+        if file_path.endswith('.py'):
+            self.print_func(f"Edit the file: {Fore.MAGENTA}{file_path}")
 
     def build_boilerplate_code(self):
         imports = self.get_imports()
@@ -183,7 +185,7 @@ class TemplateFileCreator:
             func_args = ', '.join([f"{k}: {v}" if str(v) not in ["NoneType", "None"] else str(k)
                                    for k, v in func_args.items()])
         else:
-            func_args = f"case: {self.case_type.__name__}"
+            func_args = f"case: {self.case_query.case_type.__name__}"
         return func_args
 
     def write_to_file(self, code: str):
@@ -212,7 +214,7 @@ class TemplateFileCreator:
                 else:
                     case_type_imports.append(v)
         else:
-            case_type_imports.append(self.case_type)
+            case_type_imports.append(self.case_query.case_type)
         if self.output_type is None:
             output_type_imports = [Any]
         else:
@@ -302,7 +304,7 @@ class TemplateFileCreator:
                 exec(source, scope, exec_globals)
                 user_function = exec_globals[func_name]
                 updates[func_name] = user_function
-                print_func(f"{Fore.BLUE}Loaded `{func_name}` function into user namespace.{Style.RESET_ALL}")
+                print_func(f"\n{Fore.WHITE}Loaded the following function into user namespace:\n{Fore.GREEN}{func_name}{Style.RESET_ALL}")
                 break
         if updates:
             all_code_lines = extract_function_source(file_path,