feat: add function overloading

anancarv · Mar 18, 2021 · 37783f3 · 37783f3
2 parents ab2db16 + 7491681
commit 37783f3
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/pyartifactory/objects.py b/pyartifactory/objects.py
@@ -6,7 +6,7 @@
 import logging
 import os
 from os.path import isdir, join
-from typing import List, Optional, Dict, Tuple, Union, Iterator
+from typing import List, Optional, Dict, Tuple, Union, Iterator, overload
 
 from pathlib import Path
 import requests
@@ -453,6 +453,18 @@ def get_repo(self, repo_name: str) -> AnyRepositoryResponse:
                 )
             raise ArtifactoryException from error
 
+    @overload
+    def create_repo(self, repo: LocalRepository) -> LocalRepositoryResponse:
+        ...
+
+    @overload
+    def create_repo(self, repo: VirtualRepository) -> VirtualRepositoryResponse:
+        ...
+
+    @overload
+    def create_repo(self, repo: RemoteRepository) -> RemoteRepositoryResponse:
+        ...
+
     def create_repo(self, repo: AnyRepository) -> AnyRepositoryResponse:
         """
         Creates a local, virtual or remote repository
@@ -476,6 +488,18 @@ def create_repo(self, repo: AnyRepository) -> AnyRepositoryResponse:
             logger.debug("Repository %s successfully created", repo_name)
             return self.get_repo(repo_name)
 
+    @overload
+    def update_repo(self, repo: LocalRepository) -> LocalRepositoryResponse:
+        ...
+
+    @overload
+    def update_repo(self, repo: VirtualRepository) -> VirtualRepositoryResponse:
+        ...
+
+    @overload
+    def update_repo(self, repo: RemoteRepository) -> RemoteRepositoryResponse:
+        ...
+
     def update_repo(self, repo: AnyRepository) -> AnyRepositoryResponse:
         """
         Updates a local, virtual or remote repository

diff --git a/tests/test_repositories.py b/tests/test_repositories.py
@@ -184,6 +184,7 @@ def test_create_local_repository_using_create_repo_success(mocker):
     mocker.spy(artifactory_repo, "get_repo")
     local_repo = artifactory_repo.create_repo(LOCAL_REPOSITORY)
 
+    assert isinstance(local_repo, LocalRepositoryResponse)
     assert local_repo == LOCAL_REPOSITORY_RESPONSE
 
 
@@ -582,6 +583,7 @@ def test_update_local_repository_using_update_repo_success(mocker):
     mocker.spy(artifactory_repo, "get_repo")
     updated_repo = artifactory_repo.update_repo(UPDATED_LOCAL_REPOSITORY)
 
+    assert isinstance(updated_repo, LocalRepositoryResponse)
     assert updated_repo == UPDATED_LOCAL_REPOSITORY_RESPONSE