From 7cb9c2402239dc49504cc8cbc78b92950402c064 Mon Sep 17 00:00:00 2001 From: Khaleel Al-Adhami Date: Mon, 21 Oct 2024 23:40:54 -0700 Subject: [PATCH] simplify case for when type is in the same module --- reflex/components/datadisplay/dataeditor.pyi | 23 +++++--------------- reflex/utils/pyi_generator.py | 15 ++++++++++--- 2 files changed, 17 insertions(+), 21 deletions(-) diff --git a/reflex/components/datadisplay/dataeditor.pyi b/reflex/components/datadisplay/dataeditor.pyi index e50dada20..aadd9666e 100644 --- a/reflex/components/datadisplay/dataeditor.pyi +++ b/reflex/components/datadisplay/dataeditor.pyi @@ -8,7 +8,6 @@ from typing import Any, Dict, List, Literal, Optional, Union, overload from typing_extensions import TypedDict -import reflex from reflex.base import Base from reflex.components.component import NoSSRComponent from reflex.event import EventType @@ -192,34 +191,22 @@ class DataEditor(NoSSRComponent): on_cell_context_menu: Optional[EventType[tuple[int, int]]] = None, on_cell_edited: Optional[EventType] = None, on_click: Optional[EventType[[]]] = None, - on_column_resize: Optional[ - EventType[reflex.components.datadisplay.dataeditor.GridColumn, int] - ] = None, + on_column_resize: Optional[EventType[GridColumn, int]] = None, on_context_menu: Optional[EventType[[]]] = None, - on_delete: Optional[ - EventType[reflex.components.datadisplay.dataeditor.GridSelection] - ] = None, + on_delete: Optional[EventType[GridSelection]] = None, on_double_click: Optional[EventType[[]]] = None, on_finished_editing: Optional[ - EventType[ - Union[reflex.components.datadisplay.dataeditor.GridCell, None], - list[int], - ] + EventType[Union[GridCell, None], list[int]] ] = None, on_focus: Optional[EventType[[]]] = None, on_group_header_clicked: Optional[EventType] = None, on_group_header_context_menu: Optional[ - EventType[ - int, - reflex.components.datadisplay.dataeditor.GroupHeaderClickedEventArgs, - ] + EventType[int, GroupHeaderClickedEventArgs] ] = None, on_group_header_renamed: Optional[EventType[str, str]] = None, on_header_clicked: Optional[EventType[tuple[int, int]]] = None, on_header_context_menu: Optional[EventType[tuple[int, int]]] = None, - on_header_menu_click: Optional[ - EventType[int, reflex.components.datadisplay.dataeditor.Rectangle] - ] = None, + on_header_menu_click: Optional[EventType[int, Rectangle]] = None, on_item_hovered: Optional[EventType[tuple[int, int]]] = None, on_mount: Optional[EventType[[]]] = None, on_mouse_down: Optional[EventType[[]]] = None, diff --git a/reflex/utils/pyi_generator.py b/reflex/utils/pyi_generator.py index b5f078594..1fc17341b 100644 --- a/reflex/utils/pyi_generator.py +++ b/reflex/utils/pyi_generator.py @@ -375,12 +375,13 @@ def _extract_class_props_as_ast_nodes( return kwargs -def type_to_ast(typ) -> ast.AST: +def type_to_ast(typ, cls: type) -> ast.AST: """Converts any type annotation into its AST representation. Handles nested generic types, unions, etc. Args: typ: The type annotation to convert. + cls: The class where the type annotation is used. Returns: The AST representation of the type annotation. @@ -394,6 +395,14 @@ def type_to_ast(typ) -> ast.AST: if origin is None: if hasattr(typ, "__name__"): if typ.__module__.startswith("reflex."): + typ_parts = typ.__module__.split(".") + cls_parts = cls.__module__.split(".") + + zipped = list(zip(typ_parts, cls_parts, strict=False)) + + if all(a == b for a, b in zipped) and len(typ_parts) == len(cls_parts): + return ast.Name(id=typ.__name__) + return ast.Name(id=typ.__module__ + "." + typ.__name__) return ast.Name(id=typ.__name__) elif hasattr(typ, "_name"): @@ -411,7 +420,7 @@ def type_to_ast(typ) -> ast.AST: return ast.Name(id=base_name) # Convert all type arguments recursively - arg_nodes = [type_to_ast(arg) for arg in args] + arg_nodes = [type_to_ast(arg, cls) for arg in args] # Special case for single-argument types (like List[T] or Optional[T]) if len(arg_nodes) == 1: @@ -492,7 +501,7 @@ def _generate_component_create_functiondef( ] # Convert each argument type to its AST representation - type_args = [type_to_ast(arg) for arg in arguments_without_var] + type_args = [type_to_ast(arg, cls=clz) for arg in arguments_without_var] # Join the type arguments with commas for EventType args_str = ", ".join(ast.unparse(arg) for arg in type_args)