Added type validation for websocket messages (#20)

Author: idling-mind

src/refast/components/shadcn/charts/funnel.py

@@ -111,8 +111,7 @@ def __init__(
         # Auto-inject fill colors by index when data items have no fill
         if data and not any("fill" in item for item in data):
             self.data = [
-                {**item, "fill": f"hsl(var(--chart-{(i % 8) + 1}))"}
-                for i, item in enumerate(data)
+                {**item, "fill": f"hsl(var(--chart-{(i % 8) + 1}))"} for i, item in enumerate(data)
             ]
         else:
             self.data = list(data)

src/refast/components/shadcn/charts/pie.py

@@ -131,8 +131,7 @@ def __init__(
         # Auto-inject fill colors by index when data items have no fill
         if data and not any("fill" in item for item in data):
             self.data = [
-                {**item, "fill": f"hsl(var(--chart-{(i % 8) + 1}))"}
-                for i, item in enumerate(data)
+                {**item, "fill": f"hsl(var(--chart-{(i % 8) + 1}))"} for i, item in enumerate(data)
             ]
         else:
             self.data = list(data)

src/refast/components/shadcn/charts/radial.py

@@ -50,10 +50,7 @@ def __init__(
             ]
         # Remap name_key → "name" so Recharts legend/tooltip can read entry.name
         if name_key and name_key != "name":
-            processed = [
-                {**item, "name": item[name_key]}
-                for item in processed
-            ]
+            processed = [{**item, "name": item[name_key]} for item in processed]
         self.data = processed
         self.margin = margin or {"top": 0, "right": 0, "left": 0, "bottom": 0}
         self.cx = cx

src/refast/components/shadcn/charts/scatter.py

@@ -187,9 +187,7 @@ def render(self) -> dict[str, Any]:
         # erases color/fill from scatter tooltip dimension entries, so the only
         # reliable source is the data-point payload itself.
         data_with_fill = (
-            [{**item, "fill": self.fill} for item in self.data]
-            if self.data is not None
-            else None
+            [{**item, "fill": self.fill} for item in self.data] if self.data is not None else None
         )
         return {
             "type": self.component_type,

src/refast/components/shadcn/typography.py

@@ -303,8 +303,8 @@ def _parse_custom_tags(self) -> str:
 
         # Innermost first parsing loop
         # We search for self-closing or container tags matching <[A-Z]
-        self_closing_rx = re.compile(r'<([A-Z][a-zA-Z0-9_]*)(?:\s+([^>]*?))?\s*\/>')
-        container_rx = re.compile(r'<([A-Z][a-zA-Z0-9_]*)(?:\s+([^>]*?))?\s*>(.*?)</\1>', re.DOTALL)
+        self_closing_rx = re.compile(r"<([A-Z][a-zA-Z0-9_]*)(?:\s+([^>]*?))?\s*\/>")
+        container_rx = re.compile(r"<([A-Z][a-zA-Z0-9_]*)(?:\s+([^>]*?))?\s*>(.*?)</\1>", re.DOTALL)
 
         def parse_attributes(attrs_str: str) -> dict[str, Any]:
             if not attrs_str:
@@ -316,7 +316,7 @@ def parse_attributes(attrs_str: str) -> dict[str, Any]:
             for match in attr_rx.finditer(attrs_str):
                 name = match.group(1)
                 val_group = match.group(0)
-                if '=' in val_group:
+                if "=" in val_group:
                     val = match.group(2) or match.group(3) or match.group(4) or ""
                 else:
                     val = True
@@ -338,16 +338,16 @@ def parse_attributes(attrs_str: str) -> dict[str, Any]:
                         comp_id = f"{tag_name}_{tag_counts[tag_name]}"
                         self.custom_components[comp_id] = instance
                         replacement = f"![{tag_name}](/refast-component/{comp_id})"
-                        content = content[:match.start()] + replacement + content[match.end():]
+                        content = content[: match.start()] + replacement + content[match.end() :]
                     except Exception:
                         # Validation failed (or other error), leave tag text
                         # as is (mark temporarily)
                         failed_marker = f"<__FAILED_SELF_{tag_name} {attrs_str or ''} />"
-                        content = content[:match.start()] + failed_marker + content[match.end():]
+                        content = content[: match.start()] + failed_marker + content[match.end() :]
                 else:
                     # Tag name not registered in custom_tags, leave it as is
                     failed_marker = f"<__FAILED_SELF_{tag_name} {attrs_str or ''} />"
-                    content = content[:match.start()] + failed_marker + content[match.end():]
+                    content = content[: match.start()] + failed_marker + content[match.end() :]
                 continue
 
             # 2. Try to find the first container tag with no nested un-processed tags.
@@ -358,7 +358,7 @@ def parse_attributes(attrs_str: str) -> dict[str, Any]:
                 # If inner_text contains '<[A-Z]', there is an inner un-processed tag. Skip for now.
                 # Note: we exclude '<__FAILED_' markers as they are already processed/failed.
                 # So we search for '<' followed by an uppercase letter: '<[A-Z]'
-                if re.search(r'<[A-Z]', inner_text):
+                if re.search(r"<[A-Z]", inner_text):
                     continue
                 container_match = m
                 break
@@ -376,17 +376,17 @@ def parse_attributes(attrs_str: str) -> dict[str, Any]:
                     if isinstance(callable_obj, type):
                         sig = inspect.signature(callable_obj.__init__)
                         params = list(sig.parameters.keys())
-                        if 'self' in params:
-                            params.remove('self')
+                        if "self" in params:
+                            params.remove("self")
                     else:
                         sig = inspect.signature(callable_obj)
                         params = list(sig.parameters.keys())
 
                     param_name = None
-                    if 'children' in params:
-                        param_name = 'children'
-                    elif 'content' in params:
-                        param_name = 'content'
+                    if "children" in params:
+                        param_name = "children"
+                    elif "content" in params:
+                        param_name = "content"
 
                     try:
                         if param_name:
@@ -416,29 +416,29 @@ def parse_attributes(attrs_str: str) -> dict[str, Any]:
                         self.custom_components[comp_id] = instance
                         replacement = f"![{tag_name}](/refast-component/{comp_id})"
                         content = (
-                            content[:container_match.start()]
+                            content[: container_match.start()]
                             + replacement
-                            + content[container_match.end():]
+                            + content[container_match.end() :]
                         )
                     except Exception:
                         failed_marker = (
                             f"<__FAILED_CONT_{tag_name} {attrs_str or ''}>"
                             f"{inner_content}</__FAILED_CONT_{tag_name}>"
                         )
                         content = (
-                            content[:container_match.start()]
+                            content[: container_match.start()]
                             + failed_marker
-                            + content[container_match.end():]
+                            + content[container_match.end() :]
                         )
                 else:
                     failed_marker = (
                         f"<__FAILED_CONT_{tag_name} {attrs_str or ''}>"
                         f"{inner_content}</__FAILED_CONT_{tag_name}>"
                     )
                     content = (
-                        content[:container_match.start()]
+                        content[: container_match.start()]
                         + failed_marker
-                        + content[container_match.end():]
+                        + content[container_match.end() :]
                     )
                 continue
 
@@ -455,7 +455,7 @@ def parse_attributes(attrs_str: str) -> dict[str, Any]:
 
     def _traversal_children(self) -> "list[Component]":
         children = super()._traversal_children()
-        if hasattr(self, 'custom_components') and self.custom_components:
+        if hasattr(self, "custom_components") and self.custom_components:
             for comp in self.custom_components.values():
                 if isinstance(comp, Component):
                     children.append(comp)

src/refast/models/__init__.py

@@ -0,0 +1,21 @@
+"""Pydantic models for Refast."""
+
+from refast.models.messages import (
+    CallbackMessage,
+    ClientMessage,
+    EventMessage,
+    NavigateMessage,
+    StoreInitMessage,
+    StoreSyncMessage,
+    client_message_adapter,
+)
+
+__all__ = [
+    "CallbackMessage",
+    "ClientMessage",
+    "EventMessage",
+    "NavigateMessage",
+    "StoreInitMessage",
+    "StoreSyncMessage",
+    "client_message_adapter",
+]

src/refast/models/messages.py

@@ -0,0 +1,62 @@
+"""Pydantic models for incoming WebSocket messages."""
+
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
+from pydantic.alias_generators import to_camel
+
+
+class BaseMessage(BaseModel):
+    """Base model for WebSocket messages with camelCase conversion."""
+
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+    )
+
+
+class CallbackMessage(BaseMessage):
+    """Payload for callback invocation."""
+
+    type: Literal["callback"]
+    callback_id: str
+    data: dict[str, Any] = Field(default_factory=dict)
+    event_data: Any = Field(default_factory=dict)
+
+
+class StoreInitMessage(BaseMessage):
+    """Payload for store initialization."""
+
+    type: Literal["store_init"]
+    data: dict[str, Any] = Field(default_factory=dict)
+    path: str = "/"
+
+
+class NavigateMessage(BaseMessage):
+    """Payload for navigation."""
+
+    type: Literal["navigate"]
+    path: str = "/"
+
+
+class EventMessage(BaseMessage):
+    """Payload for custom event."""
+
+    type: Literal["event"]
+    event_type: str
+    data: dict[str, Any] = Field(default_factory=dict)
+
+
+class StoreSyncMessage(BaseMessage):
+    """Payload for store synchronization."""
+
+    type: Literal["store_sync"]
+    data: dict[str, Any] = Field(default_factory=dict)
+
+
+ClientMessage = Annotated[
+    CallbackMessage | StoreInitMessage | NavigateMessage | EventMessage | StoreSyncMessage,
+    Field(discriminator="type"),
+]
+
+client_message_adapter = TypeAdapter(ClientMessage)