[cards] PythonCode component (#2196)

Netflix · Jan 10, 2025 · 2fb29ae · 2fb29ae
1 parent 0631b73
commit 2fb29ae
Show file tree

Hide file tree

Showing 14 changed files with 213 additions and 60 deletions.
diff --git a/docs/cards.md b/docs/cards.md
@@ -121,7 +121,7 @@ PATH_TO_CUSTOM_HTML = 'myhtml.html'
 class CustomCard(MetaflowCard):
     type = "custom_card"
 
-    def __init__(self, options={"no_header": True}, graph=None,components=[]):
+    def __init__(self, options={"no_header": True}, graph=None, components=[], flow=None, **kwargs):
         super().__init__()
         self._no_header = True
         self._graph = graph
@@ -177,7 +177,7 @@ class CustomCard(MetaflowCard):
 
     HTML = "<html><head></head><body>{data}<body></html>"
 
-    def __init__(self, options={"no_header": True}, graph=None,components=[]):
+    def __init__(self, options={"no_header": True}, graph=None, components=[], flow=None, **kwargs):
         super().__init__()
         self._no_header = True
         self._graph = graph
@@ -276,7 +276,7 @@ class YCard(MetaflowCard):
 
     ALLOW_USER_COMPONENTS = True
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, options={}, components=[], graph=None, flow=None, **kwargs):
         self._components = components
 
     def render(self, task):

diff --git a/metaflow/cards.py b/metaflow/cards.py
@@ -8,6 +8,7 @@
     Markdown,
     VegaChart,
     ProgressBar,
+    PythonCode,
 )
 from metaflow.plugins.cards.card_modules.basic import (
     DefaultCard,

diff --git a/metaflow/plugins/cards/card_cli.py b/metaflow/plugins/cards/card_cli.py
@@ -691,10 +691,15 @@ def create(
         try:
             if options is not None:
                 mf_card = filtered_card(
-                    options=options, components=component_arr, graph=graph_dict
+                    options=options,
+                    components=component_arr,
+                    graph=graph_dict,
+                    flow=ctx.obj.flow,
                 )
             else:
-                mf_card = filtered_card(components=component_arr, graph=graph_dict)
+                mf_card = filtered_card(
+                    components=component_arr, graph=graph_dict, flow=ctx.obj.flow
+                )
         except TypeError as e:
             if render_error_card:
                 mf_card = None

diff --git a/metaflow/plugins/cards/card_modules/basic.py b/metaflow/plugins/cards/card_modules/basic.py
@@ -1,9 +1,10 @@
 import base64
 import json
 import os
-from .card import MetaflowCard, MetaflowCardComponent
+from .card import MetaflowCard, MetaflowCardComponent, with_default_component_id
 from .convert_to_native_type import TaskToDict
 import uuid
+import inspect
 
 ABS_DIR_PATH = os.path.dirname(os.path.abspath(__file__))
 RENDER_TEMPLATE_PATH = os.path.join(ABS_DIR_PATH, "base.html")
@@ -236,9 +237,28 @@ def __init__(self, data=None):
         super().__init__(title=None, subtitle=None)
         self._data = data
 
+    @with_default_component_id
     def render(self):
         datadict = super().render()
         datadict["data"] = self._data
+        if self.component_id is not None:
+            datadict["id"] = self.component_id
+        return datadict
+
+
+class PythonCodeComponent(DefaultComponent):
+
+    type = "pythonCode"
+
+    def __init__(self, data=None):
+        super().__init__(title=None, subtitle=None)
+        self._data = data
+
+    def render(self):
+        datadict = super().render()
+        datadict["data"] = self._data
+        if self.component_id is not None:
+            datadict["id"] = self.component_id
         return datadict
 
 
@@ -343,6 +363,7 @@ def __init__(
         graph=None,
         components=[],
         runtime=False,
+        flow=None,
     ):
         self._task = task
         self._only_repr = only_repr
@@ -352,6 +373,7 @@ def __init__(
         self.final_component = None
         self.page_component = None
         self.runtime = runtime
+        self.flow = flow
 
     def render(self):
         """
@@ -475,6 +497,16 @@ def render(self):
             contents=[param_component],
         ).render()
 
+        step_func = getattr(self.flow, self._task.parent.id)
+        code_table = SectionComponent(
+            title="Task Code",
+            contents=[
+                TableComponent(
+                    data=[[PythonCodeComponent(inspect.getsource(step_func)).render()]]
+                )
+            ],
+        ).render()
+
         # Don't include parameter ids + "name" in the task artifacts
         artifactlist = [
             task_data_dict["data"][k]
@@ -500,6 +532,7 @@ def render(self):
         page_contents.extend(
             [
                 metadata_table,
+                code_table,
                 parameter_table,
                 artifact_section,
             ]
@@ -546,7 +579,7 @@ class ErrorCard(MetaflowCard):
 
     RELOAD_POLICY = MetaflowCard.RELOAD_POLICY_ONCHANGE
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, options={}, components=[], graph=None, **kwargs):
         self._only_repr = True
         self._graph = None if graph is None else transform_flow_graph(graph)
         self._components = components
@@ -602,9 +635,17 @@ class DefaultCardJSON(MetaflowCard):
 
     type = "default_json"
 
-    def __init__(self, options=dict(only_repr=True), components=[], graph=None):
+    def __init__(
+        self,
+        options=dict(only_repr=True),
+        components=[],
+        graph=None,
+        flow=None,
+        **kwargs
+    ):
         self._only_repr = True
         self._graph = None if graph is None else transform_flow_graph(graph)
+        self._flow = flow
         if "only_repr" in options:
             self._only_repr = options["only_repr"]
         self._components = components
@@ -615,6 +656,7 @@ def render(self, task):
             only_repr=self._only_repr,
             graph=self._graph,
             components=self._components,
+            flow=self._flow,
         ).render()
         return json.dumps(final_component_dict)
 
@@ -629,9 +671,17 @@ class DefaultCard(MetaflowCard):
 
     type = "default"
 
-    def __init__(self, options=dict(only_repr=True), components=[], graph=None):
+    def __init__(
+        self,
+        options=dict(only_repr=True),
+        components=[],
+        graph=None,
+        flow=None,
+        **kwargs
+    ):
         self._only_repr = True
         self._graph = None if graph is None else transform_flow_graph(graph)
+        self._flow = flow
         if "only_repr" in options:
             self._only_repr = options["only_repr"]
         self._components = components
@@ -646,6 +696,7 @@ def render(self, task, runtime=False):
             graph=self._graph,
             components=self._components,
             runtime=runtime,
+            flow=self._flow,
         ).render()
         pt = self._get_mustache()
         data_dict = dict(
@@ -688,7 +739,7 @@ class BlankCard(MetaflowCard):
 
     type = "blank"
 
-    def __init__(self, options=dict(title=""), components=[], graph=None):
+    def __init__(self, options=dict(title=""), components=[], graph=None, **kwargs):
         self._graph = None if graph is None else transform_flow_graph(graph)
         self._title = ""
         if "title" in options:

diff --git a/metaflow/plugins/cards/card_modules/card.py b/metaflow/plugins/cards/card_modules/card.py
@@ -1,4 +1,5 @@
 from typing import TYPE_CHECKING
+import uuid
 
 if TYPE_CHECKING:
     import metaflow
@@ -66,7 +67,7 @@ class MetaflowCard(object):
     # FIXME document runtime_data
     runtime_data = None
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, options={}, components=[], graph=None, flow=None):
         pass
 
     def _get_mustache(self):
@@ -140,3 +141,17 @@ def render(self):
         `render` returns a string or dictionary. This class can be called on the client side to dynamically add components to the `MetaflowCard`
         """
         raise NotImplementedError()
+
+
+def create_component_id(component):
+    uuid_bit = "".join(uuid.uuid4().hex.split("-"))[:6]
+    return type(component).__name__.lower() + "_" + uuid_bit
+
+
+def with_default_component_id(func):
+    def ret_func(self, *args, **kwargs):
+        if self.component_id is None:
+            self.component_id = create_component_id(self)
+        return func(self, *args, **kwargs)
+
+    return ret_func
diff --git a/metaflow/plugins/cards/card_modules/components.py b/metaflow/plugins/cards/card_modules/components.py
@@ -1,4 +1,4 @@
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional, Union, Callable
 from .basic import (
     LogComponent,
     ErrorComponent,
@@ -7,25 +7,13 @@
     ImageComponent,
     SectionComponent,
     MarkdownComponent,
+    PythonCodeComponent,
 )
-from .card import MetaflowCardComponent
+from .card import MetaflowCardComponent, with_default_component_id
 from .convert_to_native_type import TaskToDict, _full_classname
 from .renderer_tools import render_safely
 import uuid
-
-
-def create_component_id(component):
-    uuid_bit = "".join(uuid.uuid4().hex.split("-"))[:6]
-    return type(component).__name__.lower() + "_" + uuid_bit
-
-
-def with_default_component_id(func):
-    def ret_func(self, *args, **kwargs):
-        if self.component_id is None:
-            self.component_id = create_component_id(self)
-        return func(self, *args, **kwargs)
-
-    return ret_func
+import inspect
 
 
 def _warning_with_component(component, msg):
@@ -823,3 +811,63 @@ def render(self):
         if self._chart_inside_table and "autosize" not in self._spec:
             data["spec"]["autosize"] = "fit-x"
         return data
+
+
+class PythonCode(UserComponent):
+    """
+    A component to display Python code with syntax highlighting.
+
+    Example:
+    ```python
+    @card
+    @step
+    def my_step(self):
+        # Using code_func
+        def my_function():
+            x = 1
+            y = 2
+            return x + y
+        current.card.append(
+            PythonCode(my_function)
+        )
+
+        # Using code_string
+        code = '''
+        def another_function():
+            return "Hello World"
+        '''
+        current.card.append(
+            PythonCode(code_string=code)
+        )
+    ```
+
+    Parameters
+    ----------
+    code_func : Callable[..., Any], optional, default None
+        The function whose source code should be displayed.
+    code_string : str, optional, default None
+        A string containing Python code to display.
+        Either code_func or code_string must be provided.
+    """
+
+    def __init__(
+        self,
+        code_func: Optional[Callable[..., Any]] = None,
+        code_string: Optional[str] = None,
+    ):
+        if code_func is not None:
+            self._code_string = inspect.getsource(code_func)
+        else:
+            self._code_string = code_string
+
+    @with_default_component_id
+    @render_safely
+    def render(self):
+        if self._code_string is None:
+            return ErrorComponent(
+                "`PythonCode` component requires a `code_func` or `code_string` argument. ",
+                "None provided for both",
+            ).render()
+        _code_component = PythonCodeComponent(self._code_string)
+        _code_component.component_id = self.component_id
+        return _code_component.render()
diff --git a/metaflow/plugins/cards/card_modules/main.js b/metaflow/plugins/cards/card_modules/main.js
diff --git a/metaflow/plugins/cards/card_modules/test_cards.py b/metaflow/plugins/cards/card_modules/test_cards.py
@@ -38,7 +38,7 @@ class TestEditableCard(MetaflowCard):
 
     ALLOW_USER_COMPONENTS = True
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, components=[], **kwargs):
         self._components = components
 
     def render(self, task):
@@ -52,7 +52,7 @@ class TestEditableCard2(MetaflowCard):
 
     ALLOW_USER_COMPONENTS = True
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, components=[], **kwargs):
         self._components = components
 
     def render(self, task):
@@ -64,7 +64,7 @@ class TestNonEditableCard(MetaflowCard):
 
     seperator = "$&#!!@*"
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, components=[], **kwargs):
         self._components = components
 
     def render(self, task):
@@ -193,7 +193,7 @@ class TestRefreshComponentCard(MetaflowCard):
 
     type = "test_component_refresh_card"
 
-    def __init__(self, options={}, components=[], graph=None):
+    def __init__(self, components=[], **kwargs):
         self._components = components
 
     def render(self, task) -> str:

diff --git a/metaflow/plugins/cards/component_serializer.py b/metaflow/plugins/cards/component_serializer.py
@@ -1,8 +1,8 @@
 from .card_modules import MetaflowCardComponent
+from .card_modules.card import create_component_id
 from .card_modules.basic import ErrorComponent, SectionComponent
 from .card_modules.components import (
     UserComponent,
-    create_component_id,
     StubComponent,
 )
 from .exception import ComponentOverwriteNotSupportedException