Move widget spec type into WidgetMetdadata

databrickslabs · Jun 13, 2024 · 4f9318c · 4f9318c
1 parent 960acb6
commit 4f9318c
Showing 1 changed file with 13 additions and 8 deletions.
diff --git a/src/databricks/labs/lsql/dashboards.py b/src/databricks/labs/lsql/dashboards.py
@@ -69,7 +69,6 @@ class WidgetMetadata:
     width: int = 0
     height: int = 0
     id: str = ""
-    spec_type: type[WidgetSpec] | None = None
 
     def __post_init__(self):
         if len(self.id) == 0:
@@ -79,6 +78,13 @@ def __post_init__(self):
             self.width = self.width or width
             self.height = self.height or height
 
+    @property
+    def spec_type(self) -> type[WidgetSpec]:
+        if self.path.suffix == ".md":
+            return MarkdownSpec
+        # TODO: Determine spec by reading query when supporting more specs
+        return CounterSpec
+
     @staticmethod
     def _get_width_and_height(widget_spec: type[WidgetSpec] | None) -> tuple[int, int]:
         """Get the width and height for a widget.
@@ -228,18 +234,17 @@ def _get_widgets(
     ) -> list[tuple[Widget, WidgetMetadata]]:
         widgets = []
         for widget_metadata in widgets_metadata:
-            if widget_metadata.path.suffix == ".sql":
+            widget_spec_type = widget_metadata.spec_type
+            if widget_spec_type == MarkdownSpec:
+                widget = self._get_text_widget(widget_metadata.path)
+            else:
                 dataset = datasets[widget_metadata.path.stem]
                 try:
                     widget = self._get_widget(dataset)
                 except sqlglot.ParseError as e:
-                    logger.warning(f"Parsing {dataset.query}: {e}")
+                    logger.warning(f"Parsing {widget_metadata.path}: {e}")
                     continue
-            else:
-                widget = self._get_text_widget(widget_metadata.path)
-            assert widget.spec is not None
-            widget_with_metadata = (widget, dataclasses.replace(widget_metadata, spec_type=type(widget.spec)))
-            widgets.append(widget_with_metadata)
+            widgets.append((widget, widget_metadata))
         return widgets
 
     def _get_layouts(self, widgets: list[tuple[Widget, WidgetMetadata]]) -> list[Layout]: