chore: add type-hints to gitlab/v4/objects/geo_nodes.py

JohnVillalovos · JohnVillalovos · commit 13243b752fec · 2021-11-21T22:37:56.000-08:00
diff --git a/gitlab/v4/objects/geo_nodes.py b/gitlab/v4/objects/geo_nodes.py
@@ -1,3 +1,5 @@
+from typing import Any, cast, Dict, List, TYPE_CHECKING, Union
+
 from gitlab import cli
 from gitlab import exceptions as exc
 from gitlab.base import RequiredOptional, RESTManager, RESTObject
@@ -18,7 +20,7 @@
 class GeoNode(SaveMixin, ObjectDeleteMixin, RESTObject):
     @cli.register_custom_action("GeoNode")
     @exc.on_http_error(exc.GitlabRepairError)
-    def repair(self, **kwargs):
+    def repair(self, **kwargs: Any) -> None:
         """Repair the OAuth authentication of the geo node.
 
         Args:
@@ -30,11 +32,13 @@ def repair(self, **kwargs):
         """
         path = f"/geo_nodes/{self.get_id()}/repair"
         server_data = self.manager.gitlab.http_post(path, **kwargs)
+        if TYPE_CHECKING:
+            assert isinstance(server_data, dict)
         self._update_attrs(server_data)
 
     @cli.register_custom_action("GeoNode")
     @exc.on_http_error(exc.GitlabGetError)
-    def status(self, **kwargs):
+    def status(self, **kwargs: Any) -> Dict[str, Any]:
         """Get the status of the geo node.
 
         Args:
@@ -48,7 +52,10 @@ def status(self, **kwargs):
             dict: The status of the geo node
         """
         path = f"/geo_nodes/{self.get_id()}/status"
-        return self.manager.gitlab.http_get(path, **kwargs)
+        result = self.manager.gitlab.http_get(path, **kwargs)
+        if TYPE_CHECKING:
+            assert isinstance(result, dict)
+        return result
 
 
 class GeoNodeManager(RetrieveMixin, UpdateMixin, DeleteMixin, RESTManager):
@@ -58,9 +65,12 @@ class GeoNodeManager(RetrieveMixin, UpdateMixin, DeleteMixin, RESTManager):
         optional=("enabled", "url", "files_max_capacity", "repos_max_capacity"),
     )
 
+    def get(self, id: Union[str, int], lazy: bool = False, **kwargs: Any) -> GeoNode:
+        return cast(GeoNode, super().get(id=id, lazy=lazy, **kwargs))
+
     @cli.register_custom_action("GeoNodeManager")
     @exc.on_http_error(exc.GitlabGetError)
-    def status(self, **kwargs):
+    def status(self, **kwargs: Any) -> List[Dict[str, Any]]:
         """Get the status of all the geo nodes.
 
         Args:
@@ -73,11 +83,14 @@ def status(self, **kwargs):
         Returns:
             list: The status of all the geo nodes
         """
-        return self.gitlab.http_list("/geo_nodes/status", **kwargs)
+        result = self.gitlab.http_list("/geo_nodes/status", **kwargs)
+        if TYPE_CHECKING:
+            assert isinstance(result, list)
+        return result
 
     @cli.register_custom_action("GeoNodeManager")
     @exc.on_http_error(exc.GitlabGetError)
-    def current_failures(self, **kwargs):
+    def current_failures(self, **kwargs: Any) -> List[Dict[str, Any]]:
         """Get the list of failures on the current geo node.
 
         Args:
@@ -90,4 +103,7 @@ def current_failures(self, **kwargs):
         Returns:
             list: The list of failures
         """
-        return self.gitlab.http_list("/geo_nodes/current/failures", **kwargs)
+        result = self.gitlab.http_list("/geo_nodes/current/failures", **kwargs)
+        if TYPE_CHECKING:
+            assert isinstance(result, list)
+        return result
diff --git a/pyproject.toml b/pyproject.toml
@@ -13,7 +13,6 @@ module = [
     "docs.*",
     "docs.ext.*",
     "gitlab.v4.objects.files",
-    "gitlab.v4.objects.geo_nodes",
     "gitlab.v4.objects.issues",
     "gitlab.v4.objects.jobs",
     "gitlab.v4.objects.labels",
