Support plotly.Figure

Author: T4rk1n

dash/development/_py_components_generation.py

@@ -10,7 +10,7 @@
 from dash.exceptions import NonExistentEventException
 from ._all_keywords import python_keywords
 from ._collect_nodes import collect_nodes, filter_base_nodes
-from ._py_prop_typing import get_prop_typing, shapes
+from ._py_prop_typing import get_prop_typing, shapes, custom_imports
 from .base_component import Component
 
 
@@ -150,7 +150,7 @@ def __init__(
 
         type_name = type_info.get("name")
 
-        typed = get_prop_typing(type_name, typename, prop_key, type_info)
+        typed = get_prop_typing(type_name, typename, prop_key, type_info, namespace)
 
         arg_value = f"{prop_key}: {typed} = {default_value}"
 
@@ -220,6 +220,12 @@ def generate_class_file(
     class_string = generate_class_string(
         typename, props, description, namespace, prop_reorder_exceptions, max_props
     )
+
+    custom_imp = custom_imports[namespace][typename]
+    if custom_imports:
+        import_string += "\n".join(custom_imp)
+        import_string += "\n\n"
+
     file_name = f"{typename:s}.py"
 
     file_path = os.path.join(namespace, file_name)

dash/development/_py_prop_typing.py

@@ -1,6 +1,8 @@
+import collections
 import json
 import string
 import textwrap
+
 import stringcase
 
 
@@ -10,6 +12,7 @@
     {values}
 )
 """
+custom_imports = collections.defaultdict(lambda: collections.defaultdict(list))
 
 
 def _clean_key(key):
@@ -115,10 +118,31 @@ def generate_enum(type_info, *_):
     return f"Literal[{', '.join(values)}]"
 
 
-def get_prop_typing(type_name: str, component_name: str, prop_name: str, type_info):
-    return PROP_TYPING.get(type_name, generate_any)(
+def get_prop_typing(
+    type_name: str, component_name: str, prop_name: str, type_info, namespace=None
+):
+    if namespace:
+        # Only check the namespace once
+        special = (
+            special_cases.get(namespace, {}).get(component_name, {}).get(prop_name)
+        )
+        if special:
+            return special(type_info, component_name, prop_name)
+
+    prop_type = PROP_TYPING.get(type_name, generate_any)(
         type_info, component_name, prop_name
     )
+    return prop_type
+
+
+def generate_plotly_figure(*_):
+    custom_imports["dash_core_components"]["Graph"].append(
+        "from plotly.graph_objects import Figure"
+    )
+    return "typing.Union[Figure, dict]"
+
+
+special_cases = {"dash_core_components": {"Graph": {"figure": generate_plotly_figure}}}
 
 
 PROP_TYPING = {