Merge pull request #1224 from mito-ds/mito-for-flask-py

Mito for flask py

Author: aarondr77

mitosheet/mitosheet/mito_backend.py

@@ -25,7 +25,7 @@
 from mitosheet.enterprise.mito_config import MitoConfig
 from mitosheet.errors import (MitoError, get_recent_traceback,
                               make_execution_error)
-from mitosheet.saved_analyses import write_analysis
+from mitosheet.saved_analyses import write_save_analysis_file
 from mitosheet.steps_manager import StepsManager
 from mitosheet.telemetry.telemetry_utils import (log, log_event_processed,
                                                  telemetry_turned_on)
@@ -178,7 +178,7 @@ def handle_edit_event(self, event: Dict[str, Any]) -> None:
         self.steps_manager.handle_edit_event(event)
 
         # Also, write the analysis to a file!
-        write_analysis(self.steps_manager)
+        write_save_analysis_file(self.steps_manager)
 
         # Tell the front-end to render the new sheet and new code with an empty
         # response. NOTE: in the future, we can actually send back some data
@@ -217,7 +217,7 @@ def handle_update_event(self, event: Dict[str, Any]) -> None:
                 raise make_execution_error(error_modal=False)
             raise
         # Also, write the analysis to a file!
-        write_analysis(self.steps_manager)
+        write_save_analysis_file(self.steps_manager)
 
         # Tell the front-end to render the new sheet and new code with an empty
         # response. 

mitosheet/mitosheet/mito_flask/__init__.py

@@ -0,0 +1,2 @@
+from mitosheet.mito_flask.v1.flatten_utils import (flatten_mito_backend_to_json, read_backend_state_string_to_mito_backend)
+from mitosheet.mito_flask.v1.process_event import process_mito_event

mitosheet/mitosheet/mito_flask/v1/__init__.py

@@ -0,0 +1,36 @@
+import json
+from mitosheet.mito_backend import MitoBackend
+from mitosheet.saved_analyses import get_saved_analysis_string
+
+
+def flatten_mito_backend_to_json(mito_backend: MitoBackend) -> str:
+    saved_analysis_string = get_saved_analysis_string(mito_backend.steps_manager)
+    return json.dumps({
+        'backend_state': saved_analysis_string,
+        'shared_state_variables': mito_backend.get_shared_state_variables()
+    })
+
+def read_backend_state_string_to_mito_backend(backend_state_string: str) -> MitoBackend:
+
+    mb = MitoBackend()
+    steps_manager = mb.steps_manager
+    analysis = json.loads(backend_state_string)
+
+    previous_public_interface_version = steps_manager.public_interface_version
+    previous_code_options = steps_manager.code_options
+
+    try:
+        
+        # Before we execute the steps, update to the public interface version of the saved analysis
+        steps_manager.public_interface_version = analysis['public_interface_version']
+        steps_manager.code_options = analysis['code_options']
+
+        steps_manager.execute_steps_data(new_steps_data=analysis['steps_data'])
+
+    except:
+        # If we error, reset the public interface version, and code options
+        steps_manager.public_interface_version = previous_public_interface_version
+        steps_manager.code_options = previous_code_options
+        raise  
+
+    return mb

mitosheet/mitosheet/mito_flask/v1/flatten_utils.py

@@ -0,0 +1,30 @@
+from typing import Any, Dict, Optional
+from mitosheet.mito_backend import MitoBackend
+from mitosheet.mito_flask.v1.flatten_utils import (flatten_mito_backend_to_json, read_backend_state_string_to_mito_backend)
+
+def process_mito_event(backend_state: Optional[str], mito_event: Optional[Dict[str, Any]]) -> Any:
+
+    if backend_state is None:
+        mito_backend = MitoBackend()
+        from flask import jsonify
+        return jsonify({
+            "state": flatten_mito_backend_to_json(mito_backend),
+            "response": None,
+        })
+
+    mito_backend = read_backend_state_string_to_mito_backend(backend_state)
+
+    response = None
+    def mito_send(message):
+        nonlocal response
+        response = message
+        
+    mito_backend.mito_send = mito_send
+    if mito_event:
+        mito_backend.receive_message(mito_event)
+
+    from flask import jsonify
+    return jsonify({
+        "state": flatten_mito_backend_to_json(mito_backend),
+        "response": response,
+    })

mitosheet/mitosheet/mito_flask/v1/process_event.py

@@ -5,11 +5,12 @@
 # Distributed under the terms of the GPL License.
 
 from mitosheet.saved_analyses.save_utils import (
-    read_and_upgrade_analysis, write_analysis,
+    read_and_upgrade_analysis, write_save_analysis_file,
     delete_saved_analysis, rename_saved_analysis, 
     _get_all_analysis_filenames, _delete_analyses,
     SAVED_ANALYSIS_FOLDER, read_analysis,
-    make_steps_json_obj, get_analysis_exists
+    get_steps_obj_for_saved_analysis, get_analysis_exists, 
+    get_saved_analysis_string
 )
 from mitosheet.saved_analyses.upgrade import is_prev_version
 from mitosheet.saved_analyses.upgrade import upgrade_saved_analysis_to_current_version

mitosheet/mitosheet/saved_analyses/__init__.py

@@ -118,20 +118,24 @@ def rename_saved_analysis(old_analysis_name, new_analysis_name):
     else:
         raise Exception(f'Invalid rename, with old and new analysis are {old_analysis_name} and {new_analysis_name}')
 
-
-def write_saved_analysis(analysis_path: str, steps_data: List[Dict[str, Any]], public_interface_version: int, args: List[str], code_options: CodeOptions, version: str=__version__) -> None:
+def get_saved_analysis_string(steps_manager: StepsManagerType) -> str:
+    saved_analysis_string = json.dumps({
+        'version': __version__,
+        'steps_data': get_steps_obj_for_saved_analysis(steps_manager.steps_including_skipped),
+        'public_interface_version': steps_manager.public_interface_version,
+        'args': steps_manager.original_args_raw_strings,
+        'code': steps_manager.code(),
+        'code_options': steps_manager.code_options
+    }, cls=NpEncoder)
+    return saved_analysis_string
+
+def _write_saved_analysis_file(analysis_path: str, steps_manager: StepsManagerType) -> None:
+    saved_analysis_string = get_saved_analysis_string(steps_manager)
     with open(analysis_path, 'w+') as f:
-        saved_analysis = {
-            'version': version,
-            'steps_data': steps_data,
-            'public_interface_version': public_interface_version,
-            'args': args,
-            'code_options': code_options
-        }
-        f.write(json.dumps(saved_analysis, cls=NpEncoder))
+        f.write(saved_analysis_string)
 
 
-def make_steps_json_obj(
+def get_steps_obj_for_saved_analysis(
         steps: List[Step]
     ) -> List[Dict[str, Any]]:
     """
@@ -182,7 +186,7 @@ def make_steps_json_obj(
 
     return steps_json_obj
 
-def write_analysis(steps_manager: StepsManagerType, analysis_name: Optional[str]=None) -> None:
+def write_save_analysis_file(steps_manager: StepsManagerType, analysis_name: Optional[str]=None) -> None:
     """
     Writes the analysis saved in steps_manager to
     ~/.mito/{analysis_name}. If analysis_name is none, gets the temporary
@@ -203,7 +207,4 @@ def write_analysis(steps_manager: StepsManagerType, analysis_name: Optional[str]
         analysis_name = steps_manager.analysis_name
 
     analysis_path = f'{SAVED_ANALYSIS_FOLDER}/{analysis_name}.json'
-    steps = make_steps_json_obj(steps_manager.steps_including_skipped)
-
-    # Actually write the file
-    write_saved_analysis(analysis_path, steps, steps_manager.public_interface_version, steps_manager.original_args_raw_strings, steps_manager.code_options)
+    _write_saved_analysis_file(analysis_path, steps_manager)

mitosheet/mitosheet/saved_analyses/save_utils.py

@@ -14,7 +14,7 @@
 import sys
 from mitosheet.ai.ai_utils import is_open_ai_credentials_available
 
-from mitosheet.utils import is_prev_version, is_snowflake_connector_python_installed, is_snowflake_credentials_available, is_streamlit_installed, is_dash_installed
+from mitosheet.utils import is_flask_installed, is_prev_version, is_snowflake_connector_python_installed, is_snowflake_credentials_available, is_streamlit_installed, is_dash_installed
 
 pandas_pre_1_only = pytest.mark.skipif(
     not pd.__version__.startswith('0.'), 
@@ -65,11 +65,17 @@
     not is_snowflake_connector_python_installed() or not is_snowflake_credentials_available(),
     reason='requires snowflake_connector_python package and snowflake credentials'
 )
+
 requires_streamlit = pytest.mark.skipif(
     not is_streamlit_installed(),
     reason='requires streamlit to be installed'
 )
 
+requires_flask = pytest.mark.skipif(
+    not is_flask_installed(),
+    reason='requires flask to be installed'
+)
+
 requires_dash = pytest.mark.skipif(
     not is_dash_installed(),
     reason='requires dash to be installed'

mitosheet/mitosheet/tests/decorators.py

@@ -0,0 +1,54 @@
+import json
+import os
+import pandas as pd
+import pytest
+
+from mitosheet.tests.test_utils import create_mito_wrapper
+from mitosheet.tests.decorators import requires_flask
+from mitosheet.mito_flask.v1.flatten_utils import (flatten_mito_backend_to_json,
+                                                read_backend_state_string_to_mito_backend)
+
+TEST_DFS = [
+    (pd.DataFrame({"a": [1, 2], "b": ['abc', 'def'], "c": [True, False], "d": [1.0, 2.0]})),
+]
+
+
+@pytest.mark.parametrize("df", TEST_DFS)
+@requires_flask
+def test_convert_to_string(df):
+    mito = create_mito_wrapper()
+
+    df.to_csv("test.csv", index=False)
+    mito.simple_import(["test.csv"])
+
+    mb = mito.mito_backend
+    s = flatten_mito_backend_to_json(mb)
+
+    assert json.loads(s)["shared_state_variables"] == mb.get_shared_state_variables()
+
+    backend_state = json.loads(s)["backend_state"]
+    new_mb = read_backend_state_string_to_mito_backend(backend_state)
+    assert new_mb.steps_manager.dfs[0].equals(df)
+
+    os.remove("test.csv")
+
+@requires_flask
+def test_process_edit_convert_to_string_then_convert_back_still_processes():
+    test_wrapper = create_mito_wrapper()
+    df = pd.DataFrame({'A': [1, 2, 3], 'B': [1, 2, 3]})
+    df.to_csv("test.csv", index=False)
+    test_wrapper.simple_import(["test.csv"])
+    test_wrapper.add_column(0, 'C')
+
+    s = flatten_mito_backend_to_json(test_wrapper.mito_backend)
+    new_mb = read_backend_state_string_to_mito_backend(json.loads(s)["backend_state"])
+
+    assert new_mb.steps_manager.dfs[0].equals(pd.DataFrame({'A': [1, 2, 3], 'B': [1, 2, 3], 'C': [0, 0, 0]}))
+
+    new_test_wrapper = create_mito_wrapper(mito_backend=new_mb)
+    new_test_wrapper.add_column(0, 'D')
+
+    assert new_mb.steps_manager.dfs[0].equals(pd.DataFrame({'A': [1, 2, 3], 'B': [1, 2, 3], 'C': [0, 0, 0], 'D': [0, 0, 0]}))
+
+    os.remove("test.csv")
+    

mitosheet/mitosheet/tests/mito_flask/__init__.py

@@ -12,7 +12,7 @@
 
 import pandas as pd
 import pytest
-from mitosheet.saved_analyses import SAVED_ANALYSIS_FOLDER, write_analysis
+from mitosheet.saved_analyses import SAVED_ANALYSIS_FOLDER, write_save_analysis_file
 from mitosheet.saved_analyses.save_utils import read_and_upgrade_analysis
 from mitosheet.step_performers.filter import FC_NUMBER_EXACTLY
 from mitosheet.tests.test_utils import (create_mito_wrapper_with_data,
@@ -32,7 +32,7 @@ def test_recover_analysis(b_value, b_formula):
     mito.set_formula(b_formula, 0, 'B', add_column=True)
     # We first write out the analysis
     analysis_name = mito.mito_backend.analysis_name
-    write_analysis(mito.mito_backend.steps_manager)
+    write_save_analysis_file(mito.mito_backend.steps_manager)
 
     df = pd.DataFrame(data={'A': [1]})
     new_mito = create_mito_wrapper(df)
@@ -51,7 +51,7 @@ def test_persist_analysis_multi_sheet(b_value, b_formula):
     mito.set_formula(b_formula, 1, 'B', add_column=True)
     # We first write out the analysis
     analysis_name = mito.mito_backend.analysis_name
-    write_analysis(mito.mito_backend.steps_manager)
+    write_save_analysis_file(mito.mito_backend.steps_manager)
 
     df1 = pd.DataFrame(data={'A': [1]})
     df2 = pd.DataFrame(data={'A': [1]})
@@ -73,7 +73,7 @@ def test_persist_rename_column():
     mito.rename_column(0, 'A', 'NEW_COLUMN')
 
     analysis_name = mito.mito_backend.analysis_name
-    write_analysis(mito.mito_backend.steps_manager)
+    write_save_analysis_file(mito.mito_backend.steps_manager)
 
     df1 = pd.DataFrame(data={'A': [1]})
 
@@ -89,7 +89,7 @@ def test_persisit_delete_column():
     mito.delete_columns(0, 'A')
 
     analysis_name = mito.mito_backend.analysis_name
-    write_analysis(mito.mito_backend.steps_manager)
+    write_save_analysis_file(mito.mito_backend.steps_manager)
 
     df1 = pd.DataFrame(data={'A': [1]})
 

mitosheet/mitosheet/tests/mito_flask/test_flatten_utils.py

@@ -3,7 +3,7 @@
 import pytest
 import pandas as pd
 from mitosheet.enterprise.mito_config import MITO_CONFIG_CUSTOM_SHEET_FUNCTIONS_PATH, MITO_CONFIG_VERSION
-from mitosheet.saved_analyses.save_utils import write_analysis
+from mitosheet.saved_analyses.save_utils import write_save_analysis_file
 from mitosheet.tests.test_mito_config import delete_all_mito_config_environment_variables
 
 from mitosheet.tests.test_utils import create_mito_wrapper
@@ -38,7 +38,7 @@ def ADD1(col):
     assert mito.dfs[0].equals(pd.DataFrame({'A': [1, 2, 3], 'B': [2, 3, 4]}))
 
     analysis_name = mito.mito_backend.analysis_name
-    write_analysis(mito.mito_backend.steps_manager)
+    write_save_analysis_file(mito.mito_backend.steps_manager)
 
     new_mito = create_mito_wrapper(df)
     new_mito.replay_analysis(analysis_name)
@@ -55,7 +55,7 @@ def ADD1(col):
     assert mito.dfs[0].equals(pd.DataFrame({'A': [1, 2, 3], 'B': [2, 3, 4]}))
 
     analysis_name = mito.mito_backend.analysis_name
-    write_analysis(mito.mito_backend.steps_manager)
+    write_save_analysis_file(mito.mito_backend.steps_manager)
 
     new_mito = create_mito_wrapper(df, sheet_functions=[ADD1])
     new_mito.replay_analysis(analysis_name)

mitosheet/mitosheet/tests/saved_analyses/test_save_utils.py

@@ -1691,12 +1691,14 @@ def create_mito_wrapper(
         sheet_functions: Optional[List[Callable]]=None,
         importers: Optional[List[Callable]]=None,
         editors: Optional[List[Callable]]=None,
+        mito_backend: Optional[MitoBackend]=None
     ) -> MitoWidgetTestWrapper:
     """
     Creates a MitoWidgetTestWrapper with a mito instance with the given
     data frames.
     """
-    mito_backend = get_mito_backend(*args, user_defined_functions=sheet_functions, user_defined_importers=importers, user_defined_editors=editors)
+    if mito_backend is None:
+        mito_backend = get_mito_backend(*args, user_defined_functions=sheet_functions, user_defined_importers=importers, user_defined_editors=editors)
     test_wrapper =  MitoWidgetTestWrapper(mito_backend)
 
     if arg_names is not None:

mitosheet/mitosheet/tests/test_user_functions.py

@@ -7,7 +7,7 @@
 Saves the analysis with the passed name.
 """
 
-from mitosheet.saved_analyses import write_analysis
+from mitosheet.saved_analyses import write_save_analysis_file
 
 SAVE_ANALYSIS_UPDATE_EVENT = 'save_analysis_update'
 SAVE_ANALYSIS_UPDATE_PARAMS = ['analysis_name']
@@ -19,7 +19,7 @@ def execute_save_analysis_update(
     """
     Saves the analysis with the passed name
     """
-    write_analysis(steps_manager, analysis_name)
+    write_save_analysis_file(steps_manager, analysis_name)
 
 
 SAVE_ANALYSIS_UPDATE = {

mitosheet/mitosheet/tests/test_utils.py

@@ -528,6 +528,14 @@ def is_streamlit_installed() -> bool:
         return True
     except ImportError:
         return False
+        return False
+    
+def is_flask_installed() -> bool:
+    try:
+        import flask
+        return True
+    except ImportError:
+        return False
 
 def is_dash_installed() -> bool:
     try:

mitosheet/mitosheet/updates/save_analysis.py

@@ -145,7 +145,8 @@ def get_data_files_from_data_files_spec(
             # snowflake-connect-python requires at least Python 3.7
             'snowflake-connector-python[pandas]; python_version>="3.7"',
             'streamlit>=1.24',
-            'dash>=2.9'
+            'dash>=2.9',
+            "flask"
         ]
     },
     zip_safe                = False,