mage-ai 0.9.70__py3-none-any.whl → 0.9.71__py3-none-any.whl

mage_ai/server/utils/output_display.py

@@ -1,6 +1,9 @@
+from __future__ import annotations
+
 import json
+import os
 import re
-from typing import Dict, List
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from mage_ai.data_preparation.models.block.dynamic.utils import (
     has_reduce_output_from_upstreams,
@@ -51,7 +54,6 @@ def find_index_of_last_expression_lines(code_lines: List[str]) -> int:
         or paranthesis_close > paranthesis_open
         or square_brackets_close > square_brackets_open
     ):
-
         starting_index -= 1
 
         brackets_close += code_lines[starting_index].count('}')
@@ -84,12 +86,16 @@ def get_content_inside_triple_quotes(parts):
         if re.search(r'[\w]+[ ]*=[ ]*[f]*"""', first_line):
             variable = first_line.split('=')[0].strip()
 
-        return '\n'.join(parts[start_index + 1:-1]).replace('\"', '\\"'), variable
+        return '\n'.join(
+            parts[start_index + 1 : -1],
+        ).replace('"', '\\"'), variable  # ruff: ignore=E203
 
     return None, None
 
 
-def add_internal_output_info(block, code: str) -> str:
+def add_internal_output_info(
+    block, code: str, extension_uuid: Optional[str] = None, widget: bool = False
+) -> str:
     if code.startswith('%%sql') or code.startswith('%%bash') or len(code) == 0:
         return code
     code_lines = remove_comments(code.split('\n'))
@@ -139,147 +145,37 @@ def add_internal_output_info(block, code: str) -> str:
         is_dynamic = is_dynamic_block(block) if block else False
         is_dynamic_child = is_dynamic_block_child(block) if block else False
 
+        pipeline_uuid = block.pipeline.uuid if block.pipeline else None
+        repo_path = block.pipeline.repo_path if block.pipeline else None
+        block_uuid = block.uuid
+
+        replacements = [
+            ('DATAFRAME_ANALYSIS_MAX_COLUMNS', DATAFRAME_ANALYSIS_MAX_COLUMNS),
+            ('DATAFRAME_SAMPLE_COUNT_PREVIEW', DATAFRAME_SAMPLE_COUNT_PREVIEW),
+            ('block_uuid', f"'{block_uuid}'"),
+            ('has_reduce_output', has_reduce_output),
+            ('is_dynamic', is_dynamic),
+            ('is_dynamic_child', is_dynamic_child),
+            ('is_print_statement', is_print_statement),
+            ('last_line', last_line),
+            ('pipeline_uuid', f"'{pipeline_uuid}'"),
+            ('repo_path', f"'{repo_path}'"),
+            ('widget', widget),
+        ]
+        replacements.append((
+            'extension_uuid',
+            f"'{extension_uuid}'" if extension_uuid else 'None',
+        ))
+
+        custom_output_function_code = __interpolate_code_content(
+            'custom_output.py',
+            replacements,
+        )
+
         internal_output = f"""
 # Post processing code below (source: output_display.py)
 
-
-def __custom_output():
-    import json
-    import warnings
-    from datetime import datetime
-
-    import pandas as pd
-    import polars as pl
-    import simplejson
-
-    from mage_ai.data_preparation.models.block.dynamic.utils import transform_output_for_display
-    from mage_ai.data_preparation.models.block.dynamic.utils import (
-        combine_transformed_output_for_multi_output,
-        transform_output_for_display_dynamic_child,
-        transform_output_for_display_reduce_output,
-    )
-    from mage_ai.shared.environments import is_debug
-    from mage_ai.shared.parsers import encode_complex, sample_output
-
-
-    if pd.__version__ < '1.5.0':
-        from pandas.core.common import SettingWithCopyWarning
-    else:
-        from pandas.errors import SettingWithCopyWarning
-
-    warnings.simplefilter(action='ignore', category=SettingWithCopyWarning)
-
-    _internal_output_return = {last_line}
-
-    # Dynamic block child logic always takes precedence over dynamic block logic
-    if bool({is_dynamic_child}):
-        output_transformed = []
-
-        if _internal_output_return and isinstance(_internal_output_return, list):
-            for output in _internal_output_return:
-                output_tf = transform_output_for_display_dynamic_child(
-                    output,
-                    is_dynamic=bool({is_dynamic}),
-                    sample_count={DATAFRAME_ANALYSIS_MAX_COLUMNS},
-                )
-                output_transformed.append(output_tf)
-
-        output_transformed = output_transformed[:{DATAFRAME_SAMPLE_COUNT_PREVIEW}]
-
-        try:
-            _json_string = simplejson.dumps(
-                combine_transformed_output_for_multi_output(output_transformed),
-                default=encode_complex,
-                ignore_nan=True,
-            )
-
-            return print(f'[__internal_output__]{{_json_string}}')
-        except Exception as err:
-            print(type(_internal_output_return))
-            print(type(output_transformed))
-            raise err
-    elif bool({is_dynamic}):
-        _json_string = simplejson.dumps(
-            transform_output_for_display(
-                _internal_output_return,
-                sample_count={DATAFRAME_ANALYSIS_MAX_COLUMNS},
-            ),
-            default=encode_complex,
-            ignore_nan=True,
-        )
-        return print(f'[__internal_output__]{{_json_string}}')
-    elif bool({has_reduce_output}) and bool({is_dynamic_child}):
-        _json_string = simplejson.dumps(
-            transform_output_for_display_reduce_output(
-                _internal_output_return,
-                sample_count={DATAFRAME_ANALYSIS_MAX_COLUMNS},
-            ),
-            default=encode_complex,
-            ignore_nan=True,
-        )
-        return print(f'[__internal_output__]{{_json_string}}')
-    elif isinstance(_internal_output_return, pd.DataFrame) and (
-        type(_internal_output_return).__module__ != 'geopandas.geodataframe'
-    ):
-        _sample = _internal_output_return.iloc[:{DATAFRAME_SAMPLE_COUNT_PREVIEW}]
-        _columns = _sample.columns.tolist()[:{DATAFRAME_ANALYSIS_MAX_COLUMNS}]
-        for col in _columns:
-            try:
-                _sample[col] = _sample[col].fillna('')
-            except Exception:
-                pass
-        _rows = simplejson.loads(_sample[_columns].to_json(
-            date_format='iso',
-            default_handler=str,
-            orient='split',
-        ))['data']
-        _shape = _internal_output_return.shape
-        _index = _sample.index.tolist()
-
-        _json_string = simplejson.dumps(
-            dict(
-                data=dict(
-                    columns=_columns,
-                    index=_index,
-                    rows=_rows,
-                    shape=_shape,
-                ),
-                type='table',
-            ),
-            default=encode_complex,
-            ignore_nan=True,
-        )
-        return print(f'[__internal_output__]{{_json_string}}')
-    elif isinstance(_internal_output_return, pl.DataFrame):
-        return print(_internal_output_return)
-    elif type(_internal_output_return).__module__ == 'pyspark.sql.dataframe':
-        _sample = _internal_output_return.limit({DATAFRAME_SAMPLE_COUNT_PREVIEW}).toPandas()
-        _columns = _sample.columns.tolist()[:40]
-        _rows = _sample.to_numpy().tolist()
-        _shape = [_internal_output_return.count(), len(_sample.columns.tolist())]
-        _index = _sample.index.tolist()
-
-        _json_string = simplejson.dumps(
-            dict(
-                data=dict(
-                    columns=_columns,
-                    index=_index,
-                    rows=_rows,
-                    shape=_shape,
-                ),
-                type='table',
-            ),
-            default=encode_complex,
-            ignore_nan=True,
-        )
-        return print(f'[__internal_output__]{{_json_string}}')
-    elif not {is_print_statement}:
-        output, sampled = sample_output(encode_complex(_internal_output_return))
-        if sampled:
-            print('Sampled output is provided here for preview.')
-        return output
-
-    return
+{custom_output_function_code}
 
 __custom_output()
 """
@@ -291,37 +187,48 @@ __custom_output()
         return custom_code
 
 
+def __interpolate_code_content(
+    file_name: str,
+    replacements: List[Tuple[str, Union[bool, float, int, str, Dict[str, Any]]]],
+) -> str:
+    current_directory = os.path.dirname(__file__)
+    path_to_file = os.path.join(current_directory, file_name)
+
+    with open(path_to_file, 'r') as file:
+        content = file.read()
+
+        for placeholder, replacement in replacements:
+            placeholder_pattern = f"'{{{placeholder}}}'"
+            content = re.sub(placeholder_pattern, str(replacement), content)
+
+        return content
+
+
 def add_execution_code(
     pipeline_uuid: str,
     block_uuid: str,
     code: str,
     global_vars,
-    block_type: BlockType = None,
-    execution_uuid: str = None,
-    extension_uuid: str = None,
-    kernel_name: str = None,
+    repo_path: str,
+    block_type: Optional[BlockType] = None,
+    execution_uuid: Optional[str] = None,
+    extension_uuid: Optional[str] = None,
+    kernel_name: Optional[str] = None,
     output_messages_to_logs: bool = False,
-    pipeline_config: Dict = None,
-    repo_config: Dict = None,
+    pipeline_config: Optional[Dict] = None,
+    pipeline_config_json_encoded: Optional[str] = None,
+    repo_config: Optional[Dict] = None,
+    repo_config_json_encoded: Optional[str] = None,
     run_incomplete_upstream: bool = False,
-    run_settings: Dict = None,
+    run_settings: Optional[Dict] = None,
     run_tests: bool = False,
     run_upstream: bool = False,
     update_status: bool = True,
-    upstream_blocks: List[str] = None,
-    variables: Dict = None,
+    upstream_blocks: Optional[List[str]] = None,
+    variables: Optional[Dict] = None,
     widget: bool = False,
 ) -> str:
-    escaped_code = code.replace("'''", "\"\"\"")
-
-    if execution_uuid:
-        execution_uuid = f"'{execution_uuid}'"
-
-    if extension_uuid:
-        extension_uuid = f"'{extension_uuid}'"
-    if upstream_blocks:
-        upstream_blocks = ', '.join([f"'{u}'" for u in upstream_blocks])
-        upstream_blocks = f'[{upstream_blocks}]'
+    escaped_code = code.replace("'''", '"""').replace("\\", "\\\\")
 
     run_settings_json = json.dumps(run_settings or {})
 
@@ -337,147 +244,71 @@ def add_execution_code(
         else:
             if block_type in [BlockType.DATA_LOADER, BlockType.TRANSFORMER]:
                 magic_header = '%%spark -o df --maxrows 10000'
-    elif pipeline_config['type'] == 'databricks':
-        spark_session_init = '''
+    elif pipeline_config and pipeline_config['type'] == 'databricks':
+        spark_session_init = """
 from pyspark.sql import SparkSession
 spark = SparkSession.builder.getOrCreate()
-'''
-
-    return f"""{magic_header}
-from mage_ai.data_preparation.models.block.dynamic.utils import build_combinations_for_dynamic_child
-from mage_ai.data_preparation.models.block.dynamic.utils import has_reduce_output_from_upstreams
-from mage_ai.data_preparation.models.block.dynamic.utils import is_dynamic_block
-from mage_ai.data_preparation.models.block.dynamic.utils import is_dynamic_block_child
-from mage_ai.data_preparation.models.pipeline import Pipeline
-from mage_ai.settings.repo import get_repo_path
-from mage_ai.orchestration.db import db_connection
-from mage_ai.shared.array import find
-from mage_ai.shared.hash import merge_dict
-import datetime
-import json
-import logging
-import pandas as pd
-
-
-db_connection.start_session()
-{spark_session_init}
-
-if 'context' not in globals():
-    context = dict()
-
-def execute_custom_code():
-    block_uuid=\'{block_uuid}\'
-    run_incomplete_upstream={str(run_incomplete_upstream)}
-    run_upstream={str(run_upstream)}
-    pipeline = Pipeline(
-        uuid=\'{pipeline_uuid}\',
-        config={pipeline_config},
-        repo_config={repo_config},
-    )
-
-    block = pipeline.get_block(block_uuid, extension_uuid={extension_uuid}, widget={widget})
-
-    upstream_blocks = {upstream_blocks}
-    if upstream_blocks and len(upstream_blocks) >= 1:
-        blocks = pipeline.get_blocks({upstream_blocks})
-        block.upstream_blocks = blocks
+"""
 
-    code = r\'\'\'
+    replacements = [
+        ('DATAFRAME_SAMPLE_COUNT_PREVIEW', DATAFRAME_SAMPLE_COUNT_PREVIEW),
+        ('block_uuid', f"'{block_uuid}'"),
+        (
+            'escaped_code',
+            f"""r\'\'\'
 {escaped_code}
-    \'\'\'
-
-    global_vars = merge_dict({global_vars} or dict(), pipeline.variables or dict())
-
-    if {variables}:
-        global_vars = merge_dict(global_vars, {variables})
-
-    if pipeline.run_pipeline_in_one_process:
-        # Use shared context for blocks
-        global_vars['context'] = context
+\'\'\'""",
+        ),
+        ('global_vars', global_vars),
+        ('output_messages_to_logs', output_messages_to_logs),
+        ('pipeline_config_json_encoded', f"'{pipeline_config_json_encoded}'"),
+        ('pipeline_uuid', f"'{pipeline_uuid}'"),
+        ('repo_config_json_encoded', f"'{repo_config_json_encoded}'"),
+        ('repo_path', f"'{repo_path}'"),
+        ('run_incomplete_upstream', run_incomplete_upstream),
+        ('run_settings_json', f"'{run_settings_json}'"),
+        ('run_tests', run_tests),
+        ('run_upstream', run_upstream),
+        ('spark', 'spark'),
+        ('spark_session_init', f"'{spark_session_init}'"),
+        ('update_status', update_status),
+        ('variables', variables),
+        ('widget', widget),
+    ]
 
-    try:
-        global_vars[\'spark\'] = spark
-    except Exception:
-        pass
-
-    if run_incomplete_upstream or run_upstream:
-        block.run_upstream_blocks(
-            from_notebook=True,
-            global_vars=global_vars,
-            incomplete_only=run_incomplete_upstream,
-        )
+    if execution_uuid:
+        replacements.append(('execution_uuid', f"'{execution_uuid}'"))
 
-    logger = logging.getLogger('{block_uuid}_test')
-    logger.setLevel('INFO')
-    if 'logger' not in global_vars:
-        global_vars['logger'] = logger
+    replacements.append(('extension_uuid', f"'{extension_uuid}'" if extension_uuid else 'None'))
 
-    block_output = dict(output=[])
-    options = dict(
-        custom_code=code,
-        execution_uuid={execution_uuid},
-        from_notebook=True,
-        global_vars=global_vars,
-        logger=logger,
-        output_messages_to_logs={output_messages_to_logs},
-        run_settings=json.loads('{run_settings_json}'),
-        update_status={update_status},
-    )
+    if upstream_blocks:
+        upstream_blocks = ', '.join([f"'{u}'" for u in upstream_blocks])
+        upstream_blocks = f'[{upstream_blocks}]'
+        replacements.append(('upstream_blocks', upstream_blocks))
 
-    has_reduce_output = has_reduce_output_from_upstreams(block)
-    is_dynamic_child = is_dynamic_block_child(block)
-
-    if is_dynamic_child:
-        outputs = []
-        settings = build_combinations_for_dynamic_child(block, **options)
-        for dynamic_block_index, config in enumerate(settings):
-            output_dict = block.execute_with_callback(**merge_dict(options, config))
-            if output_dict and output_dict.get('output'):
-                outputs.append(output_dict.get('output'))
-
-            if {run_tests}:
-                block.run_tests(
-                    custom_code=code,
-                    dynamic_block_index=dynamic_block_index,
-                    from_notebook=True,
-                    logger=logger,
-                    global_vars=global_vars,
-                    update_tests=False,
-                )
-
-        block_output['output'] = outputs
-    else:
-        block_output = block.execute_with_callback(**options)
+    code_content = __interpolate_code_content('execute_custom_code.py', replacements)
 
-        if {run_tests}:
-            block.run_tests(
-                custom_code=code,
-                from_notebook=True,
-                logger=logger,
-                global_vars=global_vars,
-                update_tests=False,
-            )
+    return f"""{magic_header}
 
-    output = block_output['output'] or []
+import datetime
 
-    if {widget} or is_dynamic_block(block) or is_dynamic_child:
-        return output
-    else:
-        return find(lambda val: val is not None, output)
+{code_content}
 
-df = execute_custom_code()
-    """
+df = await execute_custom_code()
+"""
 
 
 def get_block_output_process_code(
     pipeline_uuid: str,
     block_uuid: str,
-    block_type: BlockType = None,
-    kernel_name: str = None,
-
+    repo_path: str,
+    block_type: Optional[BlockType] = None,
+    kernel_name: Optional[str] = None,
 ):
-    if kernel_name != KernelName.PYSPARK or \
-            block_type not in [BlockType.DATA_LOADER, BlockType.TRANSFORMER]:
+    if kernel_name != KernelName.PYSPARK or block_type not in [
+        BlockType.DATA_LOADER,
+        BlockType.TRANSFORMER,
+    ]:
         return None
     return f"""%%local
 from mage_ai.data_preparation.models.constants import BlockStatus
@@ -488,6 +319,7 @@ import pandas
 block_uuid=\'{block_uuid}\'
 pipeline = Pipeline(
     uuid=\'{pipeline_uuid}\',
+    repo_path=\'{repo_path}\',
 )
 block = pipeline.get_block(block_uuid)
 variable_mapping = dict(df=df)
@@ -499,19 +331,20 @@ block.update_status(BlockStatus.EXECUTED)
 
 def get_pipeline_execution_code(
     pipeline_uuid: str,
-    global_vars: Dict = None,
+    repo_path: str,
+    global_vars: Optional[Dict] = None,
     kernel_name: str = None,
-    pipeline_config: Dict = None,
-    repo_config: Dict = None,
+    pipeline_config: Optional[Dict] = None,
+    repo_config: Optional[Dict] = None,
     update_status: bool = True,
 ) -> str:
     spark_session_init = ''
     if pipeline_config['type'] == 'databricks':
-        spark_session_init = '''
+        spark_session_init = """
 from pyspark.sql import SparkSession
 import os
 spark = SparkSession.builder.master(os.getenv('SPARK_MASTER_HOST', 'local')).getOrCreate()
-'''
+"""
 
     return f"""
 from mage_ai.data_preparation.models.pipeline import Pipeline
@@ -522,6 +355,7 @@ import asyncio
 def execute_pipeline():
     pipeline = Pipeline(
         uuid=\'{pipeline_uuid}\',
+        repo_path=\'{repo_path}\',
         config={pipeline_config},
         repo_config={repo_config},
     )