plug task enqueing properly

modelgen/__main__.py

@@ -72,10 +72,24 @@ def cmd_from_schema(args):
         print("that exports DATACLASSES and ENUMS lists.", file=sys.stderr)
         sys.exit(1)
 
-    print(f"Loading schema: {schema_path}")
-    schema = load_schema(schema_path)
+    # Parse include groups
+    include = None
+    if args.include:
+        include = {g.strip() for g in args.include.split(",")}
 
-    print(f"Found {len(schema.models)} models, {len(schema.enums)} enums")
+    print(f"Loading schema: {schema_path}")
+    schema = load_schema(schema_path, include=include)
+
+    loaded = []
+    if schema.models:
+        loaded.append(f"{len(schema.models)} models")
+    if schema.enums:
+        loaded.append(f"{len(schema.enums)} enums")
+    if schema.api_models:
+        loaded.append(f"{len(schema.api_models)} api models")
+    if schema.grpc_messages:
+        loaded.append(f"{len(schema.grpc_messages)} grpc messages")
+    print(f"Found {', '.join(loaded)}")
 
     # Parse targets
     targets = [t.strip() for t in args.targets.split(",")]
@@ -237,6 +251,12 @@ def main():
         default="pydantic",
         help=f"Comma-separated output targets ({formats_str})",
     )
+    schema_parser.add_argument(
+        "--include",
+        type=str,
+        default=None,
+        help="Comma-separated model groups to include (dataclasses,enums,api,grpc). Default: all.",
+    )
     schema_parser.set_defaults(func=cmd_from_schema)
 
     # extract command

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)