Add type-hints to ensure IDEs provide right auto-completion (#504)

Signed-off-by: Sambhav Kothari <[email protected]>

Makefile

@@ -45,6 +45,7 @@ workflows-models: ## Generate the Workflows models portion of Argo Workflows
 		--disable-appending-item-suffix \
 		--disable-timestamp
 	@poetry run python scripts/models.py $(OPENAPI_SPEC_URL) workflows
+	@poetry run stubgen -o src -p hera.workflows.models && rm -rf **/__init__.pyi
 	@$(MAKE) format
 
 .PHONY: events-models
@@ -59,6 +60,7 @@ events-models: ## Generate the Events models portion of Argo Workflows
 		--disable-appending-item-suffix \
 		--disable-timestamp
 	@poetry run python scripts/models.py $(OPENAPI_SPEC_URL) events
+	@poetry run stubgen -o src -p hera.events.models && rm -rf **/__init__.pyi
 	@$(MAKE) format
 
 .PHONY: models

scripts/models.py

@@ -2,6 +2,7 @@
 
 # we want the init of `workflows.models` to have a filtered import of Workflow models
 # we can parse out the JSON using the old code and filter on Workflow objects
+import os
 import sys
 from pathlib import Path
 
@@ -129,6 +130,19 @@ def write_imports(imports: list, models_type: str, openapi_spec_url: str) -> Non
             f.write(f"from hera.{models_type}.models.io.k8s.api.core.v1 import {enum}\n")
 
 
+# Ensure that an init file is present in every folder recursively inside the hera/<models_type>/models folder
+# This is to ensure that stubgen works appropriately
+def ensure_init():
+    for models_type in model_types:
+        for root, _, files in os.walk(f"src/hera/{models_type}/models"):
+            # ignore __pycache__ folders
+            if "__pycache__" in root:
+                continue
+            if "__init__.py" not in files:
+                with open(f"{root}/__init__.py", "w") as f:
+                    f.write("")
+
+
 if __name__ == "__main__":
     models_type = get_models_type()
     openapi_spec_url = get_openapi_spec_url()

src/hera/events/models/eventsource.pyi

@@ -0,0 +1,30 @@
+from typing import Optional
+
+from hera.shared._base_model import BaseModel as BaseModel
+
+from .io.argoproj.events import v1alpha1 as v1alpha1
+from .io.k8s.apimachinery.pkg.apis.meta import v1 as v1
+
+class EventSourceDeletedResponse(BaseModel): ...
+
+class LogEntry(BaseModel):
+    event_name: Optional[str]
+    event_source_name: Optional[str]
+    event_source_type: Optional[str]
+    level: Optional[str]
+    msg: Optional[str]
+    namespace: Optional[str]
+    time: Optional[v1.Time]
+
+class CreateEventSourceRequest(BaseModel):
+    event_source: Optional[v1alpha1.EventSource]
+    namespace: Optional[str]
+
+class EventSourceWatchEvent(BaseModel):
+    object: Optional[v1alpha1.EventSource]
+    type: Optional[str]
+
+class UpdateEventSourceRequest(BaseModel):
+    event_source: Optional[v1alpha1.EventSource]
+    name: Optional[str]
+    namespace: Optional[str]

src/hera/events/models/google/protobuf.pyi

@@ -0,0 +1,7 @@
+from typing import Optional
+
+from hera.shared._base_model import BaseModel as BaseModel
+
+class Any(BaseModel):
+    type_url: Optional[str]
+    value: Optional[str]

src/hera/events/models/grpc/gateway/runtime.pyi

@@ -0,0 +1,18 @@
+from typing import List, Optional
+
+from hera.shared._base_model import BaseModel as BaseModel
+
+from ...google import protobuf as protobuf
+
+class Error(BaseModel):
+    code: Optional[int]
+    details: Optional[List[protobuf.Any]]
+    error: Optional[str]
+    message: Optional[str]
+
+class StreamError(BaseModel):
+    details: Optional[List[protobuf.Any]]
+    grpc_code: Optional[int]
+    http_code: Optional[int]
+    http_status: Optional[str]
+    message: Optional[str]