plug task enqueing properly

modelgen/loader/schema.py

@@ -65,8 +65,13 @@ class SchemaLoader:
         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():
@@ -75,34 +80,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 (TypeScript-only request/response types)
-        api_models = getattr(module, "API_MODELS", [])
-        for cls in api_models:
-            self.api_models.append(self._parse_dataclass(cls))
+        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
 
@@ -169,7 +181,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)