ripple-down-rules 0.1.2__py3-none-any.whl → 0.1.5__py3-none-any.whl

ripple_down_rules/prompt.py

@@ -2,84 +2,109 @@ import ast
 import logging
 from _ast import AST
 
-from IPython.core.interactiveshell import ExecutionInfo
 from IPython.terminal.embed import InteractiveShellEmbed
 from traitlets.config import Config
-from prompt_toolkit import PromptSession
-from prompt_toolkit.completion import WordCompleter
-from sqlalchemy.orm import DeclarativeBase as SQLTable, Session
-from typing_extensions import Any, List, Optional, Tuple, Dict, Union, Type
+from typing_extensions import List, Optional, Tuple, Dict
 
-from .datastructures import Case, PromptFor, CallableExpression, create_case, parse_string_to_expression, CaseQuery
-from .utils import capture_variable_assignment
+from .datastructures.enums import PromptFor
+from .datastructures.callable_expression import CallableExpression, parse_string_to_expression
+from .datastructures.dataclasses import CaseQuery
+from .utils import extract_dependencies, contains_return_statement
+
+
+class CustomInteractiveShell(InteractiveShellEmbed):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.all_lines = []
+
+    def run_cell(self, raw_cell: str, **kwargs):
+        """
+        Override the run_cell method to capture return statements.
+        """
+        if contains_return_statement(raw_cell):
+            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
+        result = super().run_cell(raw_cell, **kwargs)
+        if not result.error_in_exec:
+            self.all_lines.append(raw_cell)
+        return result
 
 
 class IpythonShell:
     """
     Create an embedded Ipython shell that can be used to prompt the user for input.
     """
-    def __init__(self, variable_to_capture: str, scope: Optional[Dict] = None, header: Optional[str] = None):
+
+    def __init__(self, scope: Optional[Dict] = None, header: Optional[str] = None):
         """
         Initialize the Ipython shell with the given scope and header.
 
-        :param variable_to_capture: The variable to capture from the user input.
         :param scope: The scope to use for the shell.
         :param header: The header to display when the shell is started.
         """
-        self.variable_to_capture: str = variable_to_capture
         self.scope: Dict = scope or {}
         self.header: str = header or ">>> Embedded Ipython Shell"
         self.user_input: Optional[str] = None
-        self.shell: InteractiveShellEmbed = self._init_shell()
-        self._register_hooks()
+        self.shell: CustomInteractiveShell = self._init_shell()
+        self.all_code_lines: List[str] = []
 
     def _init_shell(self):
         """
         Initialize the Ipython shell with a custom configuration.
         """
         cfg = Config()
-        shell = InteractiveShellEmbed(config=cfg, user_ns=self.scope, banner1=self.header)
+        shell = CustomInteractiveShell(config=cfg, user_ns=self.scope, banner1=self.header)
         return shell
 
-    def _register_hooks(self):
-        """
-        Register hooks to capture specific events in the Ipython shell.
-        """
-        def capture_variable(exec_info: ExecutionInfo):
-            code = exec_info.raw_cell
-            if self.variable_to_capture not in code:
-                return
-            # use ast to find if the user is assigning a value to the variable "condition"
-            assignment = capture_variable_assignment(code, self.variable_to_capture)
-            if assignment:
-                # if the user is assigning a value to the variable "condition", update the raw_condition
-                self.user_input = assignment
-                print(f"[Captured {self.variable_to_capture}]:\n{self.user_input}")
-
-        self.shell.events.register('pre_run_cell', capture_variable)
-
     def run(self):
         """
         Run the embedded shell.
         """
         self.shell()
+        self.update_user_input_from_code_lines()
 
-
-def prompt_user_for_expression(case_query: CaseQuery, prompt_for: PromptFor,
-                               session: Optional[Session] = None) -> Tuple[str, CallableExpression]:
+    def update_user_input_from_code_lines(self):
+        """
+        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() == '':
+            self.user_input = None
+        else:
+            self.all_code_lines = extract_dependencies(self.shell.all_lines)
+            if len(self.all_code_lines) == 1:
+                if self.all_code_lines[0].strip() == '':
+                    self.user_input = None
+                else:
+                    self.user_input = self.all_code_lines[0].replace('return', '').strip()
+            else:
+                self.user_input = f"def _get_value(case):\n    "
+                self.user_input += '\n    '.join(self.all_code_lines)
+
+
+def prompt_user_for_expression(case_query: CaseQuery, prompt_for: PromptFor)\
+        -> Tuple[Optional[str], Optional[CallableExpression]]:
     """
     Prompt the user for an executable python expression to the given case query.
 
     :param case_query: The case query to prompt the user for.
     :param prompt_for: The type of information ask user about.
-    :param session: The sqlalchemy orm session.
     :return: A callable expression that takes a case and executes user expression on it.
     """
     while True:
         user_input, expression_tree = prompt_user_about_case(case_query, prompt_for)
+        if user_input is None:
+            if prompt_for == PromptFor.Conclusion:
+                print("No conclusion provided. Exiting.")
+                return None, None
+            else:
+                print("Conditions must be provided. Please try again.")
+                continue
         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, session=session)
+                                                 scope=case_query.scope)
         try:
             callable_expression(case_query.case)
             break
@@ -89,7 +114,7 @@ def prompt_user_for_expression(case_query: CaseQuery, prompt_for: PromptFor,
     return user_input, callable_expression
 
 
-def prompt_user_about_case(case_query: CaseQuery, prompt_for: PromptFor) -> Tuple[str, AST]:
+def prompt_user_about_case(case_query: CaseQuery, prompt_for: PromptFor) -> Tuple[Optional[str], Optional[AST]]:
     """
     Prompt the user for input.
 
@@ -99,27 +124,13 @@ def prompt_user_about_case(case_query: CaseQuery, prompt_for: PromptFor) -> Tupl
     """
     prompt_str = f"Give {prompt_for} for {case_query.name}"
     scope = {'case': case_query.case, **case_query.scope}
-    shell = IpythonShell(prompt_for.value, scope=scope, header=prompt_str)
-    user_input, expression_tree = prompt_user_input_and_parse_to_expression(shell=shell)
-    return user_input, expression_tree
-
-
-def get_completions(obj: Any) -> List[str]:
-    """
-    Get all completions for the object. This is used in the python prompt shell to provide completions for the user.
-
-    :param obj: The object to get completions for.
-    :return: A list of completions.
-    """
-    # Define completer with all object attributes and comparison operators
-    completions = ['==', '!=', '>', '<', '>=', '<=', 'in', 'not', 'and', 'or', 'is']
-    completions += ["isinstance(", "issubclass(", "type(", "len(", "hasattr(", "getattr(", "setattr(", "delattr("]
-    completions += list(create_case(obj).keys())
-    return completions
+    shell = IpythonShell(scope=scope, header=prompt_str)
+    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) -> Tuple[str, ast.AST]:
+                                              user_input: Optional[str] = None)\
+        -> Tuple[Optional[str], Optional[ast.AST]]:
     """
     Prompt the user for input.
 
@@ -132,23 +143,12 @@ def prompt_user_input_and_parse_to_expression(shell: Optional[IpythonShell] = No
             shell = IpythonShell() if shell is None else shell
             shell.run()
             user_input = shell.user_input
+            if user_input is None:
+                return None, None
+            print(user_input)
         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)
             user_input = None
-
-
-def get_prompt_session_for_obj(obj: Any) -> PromptSession:
-    """
-    Get a prompt session for an object.
-
-    :param obj: The object to get the prompt session for.
-    :return: The prompt session.
-    """
-    completions = get_completions(obj)
-    completer = WordCompleter(completions)
-    session = PromptSession(completer=completer)
-    return session