Fix mypy violations

Author: ErnestoLoma

pyatlan/cache/__init__.py

@@ -1,3 +0,0 @@
-from pyatlan.cache.role_cache import RoleCache
-from pyatlan.cache.classification_cache import ClassificationCache
-from pyatlan.cache.custom_metadata_cache import CustomMetadataCache

pyatlan/cache/classification_cache.py

@@ -1,10 +1,10 @@
+from typing import Optional
 
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.client.typedef import TypeDefClient
-from pyatlan.model.typedef import ClassificationDef
 from pyatlan.model.enums import AtlanTypeCategory
+from pyatlan.model.typedef import ClassificationDef
 
-from typing import Optional
 
 class ClassificationCache:
 
@@ -16,7 +16,9 @@ class ClassificationCache:
 
     @classmethod
     def _refresh_cache(cls) -> None:
-        response = TypeDefClient(AtlanClient()).get_typedefs(type=AtlanTypeCategory.CLASSIFICATION)
+        response = TypeDefClient(AtlanClient()).get_typedefs(
+            type=AtlanTypeCategory.CLASSIFICATION
+        )
         if response is not None:
             cls.cache_by_id = dict()
             cls.map_id_to_name = dict()
@@ -34,7 +36,7 @@ def get_id_for_name(cls, name: str) -> Optional[str]:
         Translate the provided human-readable classification name to its Atlan-internal ID string.
         """
         cls_id = cls.map_name_to_id.get(name)
-        if not cls_id and not name in cls.deleted_names:
+        if not cls_id and name not in cls.deleted_names:
             # If not found, refresh the cache and look again (could be stale)
             cls._refresh_cache()
             cls_id = cls.map_name_to_id.get(name)
@@ -51,7 +53,7 @@ def get_name_for_id(cls, idstr: str) -> Optional[str]:
         Translate the provided Atlan-internal classification ID string to the human-readable classification name.
         """
         cls_name = cls.map_id_to_name.get(idstr)
-        if not cls_name and not idstr in cls.deleted_ids:
+        if not cls_name and idstr not in cls.deleted_ids:
             # If not found, refresh the cache and look again (could be stale)
             cls._refresh_cache()
             cls_name = cls.map_id_to_name.get(idstr)

pyatlan/cache/custom_metadata_cache.py

@@ -1,11 +1,11 @@
+from typing import Optional
 
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.client.typedef import TypeDefClient
-from pyatlan.model.typedef import CustomMetadataDef, AttributeDef
+from pyatlan.error import LogicError, NotFoundError
 from pyatlan.model.enums import AtlanTypeCategory
-from pyatlan.error import LogicError
+from pyatlan.model.typedef import AttributeDef, CustomMetadataDef
 
-from typing import Optional
 
 class CustomMetadataCache:
 
@@ -18,7 +18,9 @@ class CustomMetadataCache:
 
     @classmethod
     def _refresh_cache(cls) -> None:
-        response = TypeDefClient(AtlanClient()).get_typedefs(type=AtlanTypeCategory.CUSTOM_METADATA)
+        response = TypeDefClient(AtlanClient()).get_typedefs(
+            type=AtlanTypeCategory.CUSTOM_METADATA
+        )
         if response is not None:
             cls.cache_by_id = dict()
             cls.map_id_to_name = dict()
@@ -44,12 +46,14 @@ def _refresh_cache(cls) -> None:
                         else:
                             if attr_name in cls.map_attr_name_to_id[type_id]:
                                 raise LogicError(
-                                    "Multiple custom attributes with exactly the same name (" + attr_name + ") found for: " + type_name,
-                                    code="ATLAN-PYTHON-500-100"
+                                    "Multiple custom attributes with exactly the same name ("
+                                    + attr_name
+                                    + ") found for: "
+                                    + type_name,
+                                    code="ATLAN-PYTHON-500-100",
                                 )
                             cls.map_attr_name_to_id[type_id][attr_name] = attr_id
 
-
     @classmethod
     def get_id_for_name(cls, name: str) -> Optional[str]:
         """
@@ -77,7 +81,9 @@ def get_name_for_id(cls, idstr: str) -> Optional[str]:
             return cls.map_id_to_name.get(idstr)
 
     @classmethod
-    def get_all_custom_attributes(cls, include_deleted: bool=False, force_refresh: bool=False) -> dict[str, list[AttributeDef]]:
+    def get_all_custom_attributes(
+        cls, include_deleted: bool = False, force_refresh: bool = False
+    ) -> dict[str, list[AttributeDef]]:
         """
         Retrieve all the custom metadata attributes. The map will be keyed by custom metadata set
         name, and the value will be a listing of all the attributes within that set (with all the details
@@ -88,6 +94,10 @@ def get_all_custom_attributes(cls, include_deleted: bool=False, force_refresh: b
         m = {}
         for type_id, cm in cls.cache_by_id.items():
             type_name = cls.get_name_for_id(type_id)
+            if not type_name:
+                raise NotFoundError(
+                    f"The type_name for {type_id} could not be found.", code="fixme"
+                )
             attribute_defs = cm.attribute_defs
             if include_deleted:
                 to_include = attribute_defs

pyatlan/cache/role_cache.py

@@ -1,9 +1,9 @@
+from typing import Optional
 
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.client.role import RoleClient
 from pyatlan.model.role import AtlanRole
 
-from typing import Optional
 
 class RoleCache:
 

pyatlan/client/atlan.py

@@ -22,13 +22,12 @@
 import logging
 import os
 
-
 import requests
+from pydantic import BaseSettings, Field, HttpUrl, PrivateAttr
 
-from pydantic import BaseSettings, HttpUrl, PrivateAttr, Field
 from pyatlan.exceptions import AtlanServiceException
-from pyatlan.utils import HTTPMethod, HTTPStatus, get_logger
 from pyatlan.model.core import AtlanObject
+from pyatlan.utils import HTTPMethod, HTTPStatus, get_logger
 
 LOGGER = get_logger()
 
@@ -85,7 +84,10 @@ def call_api(self, api, query_params=None, request_obj=None):
             return None
         elif response.status_code == api.expected_status:
             try:
-                if response.content is None or response.status_code == HTTPStatus.NO_CONTENT:
+                if (
+                    response.content is None
+                    or response.status_code == HTTPStatus.NO_CONTENT
+                ):
                     return None
                 if LOGGER.isEnabledFor(logging.DEBUG):
                     LOGGER.debug(

pyatlan/client/role.py

@@ -17,30 +17,30 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from pyatlan.client.atlan import AtlanClient
 from typing import Optional
-from pyatlan.model.role import (
-    RoleResponse
-)
-from pyatlan.utils import (
-    API,
-    ADMIN_URI,
-    HTTPMethod,
-    HTTPStatus,
-)
+
+from pyatlan.client.atlan import AtlanClient
+from pyatlan.model.role import RoleResponse
+from pyatlan.utils import ADMIN_URI, API, HTTPMethod, HTTPStatus
 
 
 class RoleClient:
     ROLE_API = ADMIN_URI + "roles"
-    
+
     # Role APIs
     GET_ROLES = API(ROLE_API, HTTPMethod.GET, HTTPStatus.OK)
 
     def __init__(self, client: AtlanClient):
         self.client = client
 
-
-    def get_roles(self, limit: int, filter: Optional[str]=None, sort: Optional[str]=None, count: bool=True, offset: int=0) -> RoleResponse:
+    def get_roles(
+        self,
+        limit: int,
+        filter: Optional[str] = None,
+        sort: Optional[str] = None,
+        count: bool = True,
+        offset: int = 0,
+    ) -> RoleResponse:
         if filter is None:
             filter = ""
         if sort is None:
@@ -50,13 +50,14 @@ def get_roles(self, limit: int, filter: Optional[str]=None, sort: Optional[str]=
             "sort": sort,
             "count": count,
             "offset": offset,
-            "limit": limit
+            "limit": limit,
         }
-        raw_json = self.client.call_api(RoleClient.GET_ROLES.format_path_with_params(), query_params)
+        raw_json = self.client.call_api(
+            RoleClient.GET_ROLES.format_path_with_params(), query_params
+        )
         response = RoleResponse(**raw_json)
         return response
 
-
     def get_all_roles(self) -> RoleResponse:
         """
         Retrieve all roles defined in Atlan.

pyatlan/client/typedef.py

@@ -18,12 +18,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from pyatlan.model.typedef import TypeDefResponse
-from pyatlan.utils import API, BASE_URI, HTTPMethod, HTTPStatus
 from pyatlan.client.atlan import AtlanClient
-from pyatlan.model.enums import AtlanTypeCategory
-from pyatlan.model.typedef import TypeDef, ClassificationDef, CustomMetadataDef
 from pyatlan.exceptions import InvalidRequestException
+from pyatlan.model.enums import AtlanTypeCategory
+from pyatlan.model.typedef import (
+    ClassificationDef,
+    CustomMetadataDef,
+    TypeDef,
+    TypeDefResponse,
+)
+from pyatlan.utils import API, BASE_URI, HTTPMethod, HTTPStatus
 
 
 class TypeDefClient:
@@ -55,9 +59,7 @@ def get_all_typedefs(self) -> TypeDefResponse:
         return TypeDefResponse(**raw_json)
 
     def get_typedefs(self, type: AtlanTypeCategory) -> TypeDefResponse:
-        query_params = {
-            "type": type.value
-        }
+        query_params = {"type": type.value}
         raw_json = self.client.call_api(
             TypeDefClient.GET_ALL_TYPE_DEFS.format_path_with_params(),
             query_params,
@@ -74,7 +76,7 @@ def create_typedef(self, typedef: TypeDef) -> TypeDefResponse:
                 struct_defs=[],
                 entity_defs=[],
                 relationship_defs=[],
-                businessMetadataDefs=[]
+                businessMetadataDefs=[],
             )
         elif isinstance(typedef, CustomMetadataDef):
             # Set up the request payload...
@@ -84,22 +86,25 @@ def create_typedef(self, typedef: TypeDef) -> TypeDefResponse:
                 struct_defs=[],
                 entity_defs=[],
                 relationship_defs=[],
-                businessMetadataDefs=[typedef]
+                businessMetadataDefs=[typedef],
             )
         else:
             raise InvalidRequestException(
-                "Unable to create new type definitions of category: " + typedef.category.value,
-                param="category"
+                "Unable to create new type definitions of category: "
+                + typedef.category.value,
+                param="category",
             )
             # Throw an invalid request exception
         raw_json = self.client.call_api(
-            TypeDefClient.CREATE_TYPE_DEFS,
-            request_obj=payload
+            TypeDefClient.CREATE_TYPE_DEFS, request_obj=payload
         )
         return TypeDefResponse(**raw_json)
 
     def purge_typedef(self, internal_name: str) -> None:
-        self.client.call_api(TypeDefClient.DELETE_TYPE_DEF_BY_NAME.format_path_with_params(internal_name))
+        self.client.call_api(
+            TypeDefClient.DELETE_TYPE_DEF_BY_NAME.format_path_with_params(internal_name)
+        )
+
 
 if __name__ == "__main__":
     client = TypeDefClient(AtlanClient())