modelgen update

soleprint/station/tools/modelgen/loader/schema.py

@@ -5,6 +5,7 @@ Loads Python dataclasses from a schema/ folder.
 Expects the folder to have an __init__.py that exports:
 - DATACLASSES: List of dataclass types to generate
 - ENUMS: List of Enum types to include
+- API_MODELS: (optional) List of API request/response types
 - GRPC_MESSAGES: (optional) List of gRPC message types
 - GRPC_SERVICE: (optional) gRPC service definition dict
 """
@@ -60,12 +61,18 @@ class SchemaLoader:
     def __init__(self, schema_path: Path):
         self.schema_path = Path(schema_path)
         self.models: List[ModelDefinition] = []
+        self.api_models: List[ModelDefinition] = []
         self.enums: List[EnumDefinition] = []
         self.grpc_messages: List[ModelDefinition] = []
         self.grpc_service: Optional[GrpcServiceDefinition] = None
 
-    def load(self) -> "SchemaLoader":
-        """Load schema definitions from the schema folder."""
+    def load(self, include: Optional[set] = None) -> "SchemaLoader":
+        """Load schema definitions from the schema folder.
+
+        Args:
+            include: Set of groups to load (dataclasses, enums, api, grpc).
+                     None means load all groups.
+        """
         init_path = self.schema_path / "__init__.py"
 
         if not init_path.exists():
@@ -74,29 +81,41 @@ class SchemaLoader:
         # Import the schema module
         module = self._import_module(init_path)
 
+        load_all = include is None
+
         # Extract DATACLASSES
-        dataclasses = getattr(module, "DATACLASSES", [])
-        for cls in dataclasses:
-            self.models.append(self._parse_dataclass(cls))
+        if load_all or "dataclasses" in include:
+            dataclasses = getattr(module, "DATACLASSES", [])
+            for cls in dataclasses:
+                self.models.append(self._parse_dataclass(cls))
+
+        # Extract API_MODELS (request/response types)
+        if load_all or "api" in include:
+            api_models = getattr(module, "API_MODELS", [])
+            for cls in api_models:
+                self.api_models.append(self._parse_dataclass(cls))
 
         # Extract ENUMS
-        enums = getattr(module, "ENUMS", [])
-        for enum_cls in enums:
-            self.enums.append(self._parse_enum(enum_cls))
+        if load_all or "enums" in include:
+            enums = getattr(module, "ENUMS", [])
+            for enum_cls in enums:
+                self.enums.append(self._parse_enum(enum_cls))
 
         # Extract GRPC_MESSAGES (optional)
-        grpc_messages = getattr(module, "GRPC_MESSAGES", [])
-        for cls in grpc_messages:
-            self.grpc_messages.append(self._parse_dataclass(cls))
+        if load_all or "grpc" in include:
+            grpc_messages = getattr(module, "GRPC_MESSAGES", [])
+            for cls in grpc_messages:
+                self.grpc_messages.append(self._parse_dataclass(cls))
 
         # Extract GRPC_SERVICE (optional)
-        grpc_service = getattr(module, "GRPC_SERVICE", None)
-        if grpc_service:
-            self.grpc_service = GrpcServiceDefinition(
-                package=grpc_service.get("package", "service"),
-                name=grpc_service.get("name", "Service"),
-                methods=grpc_service.get("methods", []),
-            )
+        if load_all or "grpc" in include:
+            grpc_service = getattr(module, "GRPC_SERVICE", None)
+            if grpc_service:
+                self.grpc_service = GrpcServiceDefinition(
+                    package=grpc_service.get("package", "service"),
+                    name=grpc_service.get("name", "Service"),
+                    methods=grpc_service.get("methods", []),
+                )
 
         return self
 
@@ -163,7 +182,7 @@ class SchemaLoader:
         return False
 
 
-def load_schema(schema_path: str | Path) -> SchemaLoader:
+def load_schema(schema_path: str | Path, include: Optional[set] = None) -> SchemaLoader:
     """Load schema definitions from folder."""
     loader = SchemaLoader(schema_path)
-    return loader.load()
+    return loader.load(include=include)