Merge pull request #1259 from mito-ds/preconfig-conditional-format

aarondr77 · web-flow · commit a3c1b9e8e986 · 2024-03-13T15:19:05.000-04:00
Preconfig conditional format
diff --git a/mitosheet/mitosheet/code_chunks/postprocessing/set_dataframe_format_code_chunk.py b/mitosheet/mitosheet/code_chunks/postprocessing/set_dataframe_format_code_chunk.py
@@ -180,6 +180,7 @@ def get_conditional_format_code_list(state: State, sheet_index: int) -> Tuple[Op
     """Returns all the code to set the conditional formats"""
     df_name = state.df_names[sheet_index]
     df = state.dfs[sheet_index]
+
     conditional_formats = state.df_formats[sheet_index]['conditional_formats']
 
     # We get the conditional formatting results, and we filter out any columns that are 
diff --git a/mitosheet/mitosheet/mito_backend.py b/mitosheet/mitosheet/mito_backend.py
@@ -29,7 +29,7 @@
 from mitosheet.steps_manager import StepsManager
 from mitosheet.telemetry.telemetry_utils import (log, log_event_processed,
                                                  telemetry_turned_on)
-from mitosheet.types import CodeOptions, MitoTheme, ParamMetadata
+from mitosheet.types import CodeOptions, ColumnDefinintion, ColumnDefinitions, ConditionalFormat, MitoTheme, ParamMetadata
 from mitosheet.updates.replay_analysis import REPLAY_ANALYSIS_UPDATE
 from mitosheet.user.create import try_create_user_json_file
 from mitosheet.user.db import USER_JSON_PATH, get_user_field
@@ -57,6 +57,7 @@ def __init__(
             user_defined_importers: Optional[List[Callable]]=None,
             user_defined_editors: Optional[List[Callable]]=None,
             code_options: Optional[CodeOptions]=None,
+            column_definitions: Optional[List[ColumnDefinitions]]=None,
             theme: Optional[MitoTheme]=None,
         ):
         """
@@ -92,7 +93,7 @@ def __init__(
 
             if not os.path.exists(import_folder):
                 raise ValueError(f"Import folder {import_folder} does not exist. Please change the file path or create the folder.")
-        
+            
         # Set up the state container to hold private widget state
         self.steps_manager = StepsManager(
             args, 
@@ -104,6 +105,7 @@ def __init__(
             user_defined_importers=all_custom_importers,
             user_defined_editors=user_defined_editors,
             code_options=code_options,
+            column_definitions=column_definitions,
             theme=theme
         )
 
@@ -314,6 +316,7 @@ def get_mito_backend(
         user_defined_functions: Optional[List[Callable]]=None,
         user_defined_importers: Optional[List[Callable]]=None,
         user_defined_editors: Optional[List[Callable]]=None,
+        column_definitions: Optional[List[ColumnDefinitions]]=None,
     ) -> MitoBackend:
 
     # We pass in the dataframes directly to the widget
@@ -322,7 +325,8 @@ def get_mito_backend(
         analysis_to_replay=analysis_to_replay, 
         user_defined_functions=user_defined_functions, 
         user_defined_importers=user_defined_importers,
-        user_defined_editors=user_defined_editors
+        user_defined_editors=user_defined_editors,
+        column_definitions=column_definitions
     ) 
 
     return mito_backend
diff --git a/mitosheet/mitosheet/pro/conditional_formatting_utils.py b/mitosheet/mitosheet/pro/conditional_formatting_utils.py
@@ -52,6 +52,7 @@ def get_conditonal_formatting_result(
                     # are sent to the frontend
                     json_index = json.dumps(index, cls=NpEncoder)
                     formatted_result[column_id][json_index] = {'backgroundColor': backgroundColor, 'color': color}
+
         except Exception as e:
             if format_uuid not in invalid_conditional_formats:
                 invalid_conditional_formats[format_uuid] = []
diff --git a/mitosheet/mitosheet/step_performers/filter.py b/mitosheet/mitosheet/step_performers/filter.py
@@ -41,6 +41,7 @@
     FC_NUMBER_HIGHEST,
 ]
 
+
 class FilterStepPerformer(StepPerformer):
     """
     Allows you to filter a column based on some conditions and some values.
@@ -260,5 +261,5 @@ def check_filters_contain_condition_that_needs_full_df(filters: List[Union[Filte
             filter_group: FilterGroup = filter_or_group #type: ignore
             if check_filters_contain_condition_that_needs_full_df(filter_group["filters"]):
                 return True
-        
+
     return False
diff --git a/mitosheet/mitosheet/steps_manager.py b/mitosheet/mitosheet/steps_manager.py
@@ -35,10 +35,10 @@
     SnowflakeImportStepPerformer
 from mitosheet.transpiler.transpile import transpile
 from mitosheet.transpiler.transpile_utils import get_default_code_options
-from mitosheet.types import CodeOptions, MitoTheme, ParamMetadata
+from mitosheet.types import CodeOptions, ColumnDefinintion, ColumnDefinitions, MitoTheme, ParamMetadata
 from mitosheet.updates import UPDATES
 from mitosheet.user.utils import is_enterprise, is_running_test
-from mitosheet.utils import NpEncoder, dfs_to_array_for_json, get_new_id, is_default_df_names
+from mitosheet.utils import NpEncoder, dfs_to_array_for_json, get_default_df_formats, get_new_id, is_default_df_names
 from mitosheet.step_performers.utils.user_defined_function_utils import get_user_defined_importers_for_frontend, get_user_defined_editors_for_frontend
 from mitosheet.step_performers.utils.user_defined_function_utils import validate_and_wrap_sheet_functions, validate_user_defined_editors
 
@@ -187,6 +187,7 @@ def __init__(
             user_defined_importers: Optional[List[Callable]]=None,
             user_defined_editors: Optional[List[Callable]]=None,
             code_options: Optional[CodeOptions]=None,
+            column_definitions: Optional[List[ColumnDefinitions]]=None,
             theme: Optional[MitoTheme]=None,
         ):
         """
@@ -221,11 +222,11 @@ def __init__(
         # saving any data that we need to transpilate it later this
         self.preprocess_execution_data = {}
         df_names = None
-        for preprocess_step_performers in PREPROCESS_STEP_PERFORMERS:
-            args, df_names, execution_data = preprocess_step_performers.execute(args)
+        for preprocess_step_performer in PREPROCESS_STEP_PERFORMERS:
+            args, df_names, execution_data = preprocess_step_performer.execute(args)
             self.preprocess_execution_data[
-                preprocess_step_performers.preprocess_step_type()
-            ] = execution_data       
+                preprocess_step_performer.preprocess_step_type()
+            ] = execution_data    
 
         # We set the original_args_raw_strings. If we later have an args update, then these
         # are overwritten by the args update (and are actually correct). But since we don't 
@@ -255,6 +256,8 @@ def __init__(
         # The version of the public interface used by this analysis
         self.public_interface_version = 3
 
+        df_formats = get_default_df_formats(column_definitions, list(args))
+        
         # Then we initialize the analysis with just a simple initialize step
         self.steps_including_skipped: List[Step] = [
             Step(
@@ -265,7 +268,8 @@ def __init__(
                     df_names=df_names,
                     user_defined_functions=self.user_defined_functions, 
                     user_defined_importers=self.user_defined_importers,
-                    user_defined_editors=self.user_defined_editors
+                    user_defined_editors=self.user_defined_editors,
+                    df_formats=df_formats
                 ), 
                 {}
             )
diff --git a/mitosheet/mitosheet/streamlit/v1/apps/column_definitions_app.py b/mitosheet/mitosheet/streamlit/v1/apps/column_definitions_app.py
@@ -0,0 +1,37 @@
+import pandas as pd
+import streamlit as st
+from mitosheet.streamlit.v1 import spreadsheet
+
+st.set_page_config(layout="wide")
+st.title('Tesla Stock Volume Analysis')
+
+df = pd.DataFrame({
+    'A': [1, 2, 3, 4, 5, 6, 7, 8, 9],
+    'B': [1, 2, 3, 4, 5, 6, 7, 8, 9]
+})
+new_dfs, code = spreadsheet(
+    df, 
+    column_definitions=[
+        [
+            {
+                'columns': ['A', 'B'],
+                'conditional_formats': [{
+                    'filters': [{'condition': 'greater_than_or_equal', 'value': 5}], 
+                    'font_color': '#c30010', 
+                    'background_color': '#ffcbd1' 
+                }] 
+            },
+            {
+                'columns': ['A'],
+                'conditional_formats': [{
+                    'filters': [{'condition': 'less', 'value': 2}], 
+                    'font_color': '#f30010', 
+                    'background_color': '#ddcbd1' 
+                }] 
+            }
+        ],
+    ]
+)
+
+st.write(new_dfs)
+st.code(code)
diff --git a/mitosheet/mitosheet/streamlit/v1/apps/simple_dataframe_app.py b/mitosheet/mitosheet/streamlit/v1/apps/simple_dataframe_app.py
@@ -9,7 +9,7 @@
     'a': [1, 2, 3],
     'b': [4, 5, 6]
 })
-new_dfs, code = spreadsheet(df, height='1000px')
+new_dfs, code = spreadsheet(df)
 
 st.write(new_dfs)
 st.code(code)
diff --git a/mitosheet/mitosheet/streamlit/v1/spreadsheet.py b/mitosheet/mitosheet/streamlit/v1/spreadsheet.py
@@ -13,7 +13,7 @@
 
 from mitosheet.mito_backend import MitoBackend
 from mitosheet.selection_utils import get_selected_element
-from mitosheet.types import CodeOptions, ParamMetadata, ParamType
+from mitosheet.types import CodeOptions, ColumnDefinitions, ConditionalFormat, ParamMetadata, ParamType
 from mitosheet.utils import get_new_id
 
 CURRENT_MITO_ANALYSIS_VERSION = 1
@@ -269,6 +269,7 @@ def _get_mito_backend(
             _editors: Optional[List[Callable]]=None, 
             _sheet_functions: Optional[List[Callable]]=None, 
             _code_options: Optional[CodeOptions]=None,
+            _column_definitions: Optional[List[ColumnDefinitions]]=None,
             import_folder: Optional[str]=None,
             df_names: Optional[List[str]]=None,
             session_id: Optional[str]=None,
@@ -280,6 +281,7 @@ def _get_mito_backend(
             import_folder=import_folder,
             user_defined_importers=_importers, user_defined_functions=_sheet_functions, user_defined_editors=_editors,
             code_options=_code_options,
+            column_definitions = _column_definitions,
         )
 
         # Make a send function that stores the responses in a list
@@ -320,6 +322,7 @@ def spreadsheet( # type: ignore
             df_names: Optional[List[str]]=None,
             import_folder: Optional[str]=None,
             code_options: Optional[CodeOptions]=None,
+            column_definitions: Optional[List[ColumnDefinitions]]=None,
             return_type: str='default',
             height: Optional[str]=None,
             key=None
@@ -366,6 +369,7 @@ def spreadsheet( # type: ignore
             _importers=importers, 
             _editors=editors,
             _code_options=code_options,
+            _column_definitions=column_definitions,
             import_folder=import_folder,
             session_id=session_id,
             df_names=df_names, 
diff --git a/mitosheet/mitosheet/tests/api/test_get_code_snippets.py b/mitosheet/mitosheet/tests/api/test_get_code_snippets.py
@@ -1,10 +1,11 @@
 import json
+
+from pytest_httpserver import HTTPServer
 from mitosheet.enterprise.mito_config import MITO_CONFIG_CODE_SNIPPETS_URL, MITO_CONFIG_CODE_SNIPPETS_VERSION, MITO_CONFIG_VERSION
 from mitosheet.tests.test_mito_config import delete_all_mito_config_environment_variables
 from mitosheet.tests.test_utils import create_mito_wrapper
 from mitosheet.api.get_code_snippets import get_code_snippets, DEFAULT_CODE_SNIPPETS
 import os
-from pytest_httpserver import HTTPServer
 
 
 TEST_CODE_SNIPPETS = [
diff --git a/mitosheet/mitosheet/tests/streamlit/test_runnable_analysis.py b/mitosheet/mitosheet/tests/streamlit/test_runnable_analysis.py
@@ -85,7 +85,7 @@ def function(file_name_import_csv_0, file_name_import_excel_0, file_name_export_
         }
     ]
     
-    analysis = RunnableAnalysis('', None, fully_parameterized_function, param_metadata)
+    analysis = RunnableAnalysis('', None,  fully_parameterized_function, param_metadata)
     result = analysis.run()
     assert result is not None
     pd.testing.assert_frame_equal(result[0], expected_df)
@@ -154,7 +154,7 @@ def function(file_name_import_csv_0, file_name_import_excel_0, file_name_export_
     
     new_export_file_0 = str(tmp_path / 'new_export.csv')
     new_export_file_1 = str(tmp_path / 'new_export.xlsx')
-    analysis = RunnableAnalysis('', None, fully_parameterized_function, param_metadata)
+    analysis = RunnableAnalysis('', None,  fully_parameterized_function, param_metadata)
 
     # Test that get_param_metadata works
     assert analysis.get_param_metadata() == param_metadata
@@ -192,7 +192,7 @@ def function_ctqm(import_dataframe_0):
         }
     ]
 
-    analysis = RunnableAnalysis('', None, fully_parameterized_function, param_metadata)
+    analysis = RunnableAnalysis('', None,  fully_parameterized_function, param_metadata)
 
     assert analysis.get_param_metadata() == param_metadata
     assert analysis.get_param_metadata('import') == param_metadata
@@ -273,7 +273,7 @@ def function(import_dataframe_0, file_name_import_csv_0, file_name_import_excel_
     
     new_export_file_0 = str(tmp_path / 'new_export.csv')
     new_export_file_1 = str(tmp_path / 'new_export.xlsx')
-    analysis = RunnableAnalysis('', None, fully_parameterized_function, param_metadata)
+    analysis = RunnableAnalysis('', None,  fully_parameterized_function, param_metadata)
 
     # Test that get_param_metadata works
     assert analysis.get_param_metadata() == param_metadata
@@ -383,7 +383,7 @@ def function(import_dataframe_0):
         },
     ]
     
-    analysis = RunnableAnalysis('', None, fully_parameterized_function, param_metadata)
+    analysis = RunnableAnalysis('', None,  fully_parameterized_function, param_metadata)
     with pytest.raises(TypeError):
         analysis.run()
     
@@ -406,14 +406,14 @@ def function(import_dataframe_0):
         }
     ]
     
-    analysis = RunnableAnalysis('', None, fully_parameterized_function, param_metadata)
+    analysis = RunnableAnalysis('', None,  fully_parameterized_function, param_metadata)
     with pytest.raises(TypeError):
         analysis.run(testing=1)
     
 
 @requires_streamlit
 def test_to_and_from_json():
-    analysis = RunnableAnalysis('', None, simple_fn, simple_param_metadata)
+    analysis = RunnableAnalysis('', None,  simple_fn, simple_param_metadata)
     # Test that the to_json function 
     json = analysis.to_json()
     assert json is not None
@@ -445,7 +445,7 @@ def function(vari\abl"e_name{}):
         'name': 'vari\ abl"e_name{}'
     },
 ]
-    analysis = RunnableAnalysis('', None, special_characters_fn, special_characters_metadata)
+    analysis = RunnableAnalysis('', None,  special_characters_fn, special_characters_metadata)
     # Test that the to_json function 
     json = analysis.to_json()
     assert json is not None
@@ -467,7 +467,7 @@ def function():
     
     return df1
 """
-    analysis = RunnableAnalysis('', None, fully_parameterized_code, [])
+    analysis = RunnableAnalysis('', None,  fully_parameterized_code, [])
     result = analysis.run()
     pd.testing.assert_frame_equal(result, pd.DataFrame({'A': [1, 2, 3], 'B': [2, 3, 4]}))
 
diff --git a/mitosheet/mitosheet/tests/streamlit/test_streamlit.py b/mitosheet/mitosheet/tests/streamlit/test_streamlit.py
@@ -62,7 +62,7 @@
         'return type of analysis',
         [df1], {'return_type': 'analysis'},
         (
-            RunnableAnalysis('', None, '', [])
+            RunnableAnalysis('', None, '', [], 0)
         )
     )
 ]
@@ -109,3 +109,32 @@ def test_return_type_function_invalid_code_options():
         spreadsheet(df1, code_options={'as_function': False, 'call_function': False, 'function_name': 'test', 'function_params': {}}, return_type='function')
     with pytest.raises(ValueError):
         spreadsheet(df1, code_options={'as_function': True, 'call_function': True, 'function_name': 'test', 'function_params': {}}, return_type='function')
+
+@requires_streamlit   
+def test_spreadsheet_with_column_definitions():
+    f = spreadsheet(
+        df1, 
+        column_definitions=[
+            [
+                {
+                    'columns': ['A'],
+                    'conditional_formats': [{
+                        'filters': [{'condition': 'greater_than_or_equal', 'value': 5}], 
+                        'font_color': '#c30010', 
+                        'background_color': '#ffcbd1' 
+                    }] 
+                },
+                {
+                    'columns': ['A'],
+                    'conditional_formats': [{
+                        'filters': [{'condition': 'less', 'value': 2}], 
+                        'font_color': '#f30010', 
+                        'background_color': '#ddcbd1' 
+                    }] 
+                }
+            ]
+        ], 
+        code_options={'as_function': True, 'call_function': False, 'function_name': 'test', 'function_params': {}}, 
+        return_type='function'
+    )
+    assert callable(f)
diff --git a/mitosheet/mitosheet/tests/test_column_definitions.py b/mitosheet/mitosheet/tests/test_column_definitions.py
diff --git a/mitosheet/mitosheet/tests/test_mito_backend.py b/mitosheet/mitosheet/tests/test_mito_backend.py
diff --git a/mitosheet/mitosheet/transpiler/transpile.py b/mitosheet/mitosheet/transpiler/transpile.py
diff --git a/mitosheet/mitosheet/types.py b/mitosheet/mitosheet/types.py
diff --git a/mitosheet/mitosheet/utils.py b/mitosheet/mitosheet/utils.py
