mypy

Author: esoteric-ephemera

mp_api/client/contribs/_types.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import gzip
+from abc import ABCMeta, abstractmethod
 from base64 import b64decode, b64encode
 from inspect import getfullargspec
 from pathlib import Path
@@ -30,6 +31,18 @@
 j2h = Json2Html()
 
 
+class _Component(metaclass=ABCMeta):
+    """Define component which requires a from_dict method.
+
+    Mostly exists for mypy checking.
+    """
+
+    @classmethod
+    @abstractmethod
+    def from_dict(cls, dct: dict):
+        """Instantiate from a dict."""
+
+
 class MPCDict(dict):
     """Custom dictionary to display itself as HTML table with Bulma CSS."""
 
@@ -49,7 +62,7 @@ def display(self, attrs: str = f'class="table {MPCC_SETTINGS.BULMA}"'):
         return display(HTML(html)) if _in_ipython() else html
 
 
-class Table(pd.DataFrame):
+class Table(pd.DataFrame, _Component):
     """Wrapper class around pandas.DataFrame to provide display() and info()."""
 
     def display(self):
@@ -139,7 +152,7 @@ def as_dict(self):
         return dct
 
 
-class MPCStructure(PmgStructure):
+class MPCStructure(PmgStructure, _Component):
     """Wrapper class around pymatgen.Structure to provide display() and info()."""
 
     def display(self):
@@ -178,7 +191,7 @@ def from_dict(
         return super().from_dict(dct)
 
 
-class Attachment(dict):
+class Attachment(dict, _Component):
     """Wrapper class around dict to handle attachments."""
 
     def decode(self) -> bytes:
@@ -272,9 +285,7 @@ def from_file(cls, path: str | Path) -> Self:
         content = path.read_bytes()
 
         if not (
-            supported := any(
-                isinstance(kind, ftyp) for ftyp in MPCC_SETTINGS.SUPPORTED_FILETYPES
-            )
+            supported := kind in MPCC_SETTINGS.SUPPORTED_FILETYPES
         ):  # try to gzip text file
             try:
                 content = gzip.compress(content)

mp_api/client/contribs/client.py

@@ -34,8 +34,8 @@
 from bravado_core.spec import Spec, _identity, build_api_serving_url
 from bravado_core.validate import validate_object
 from bson.objectid import ObjectId
-from cachetools import LRUCache, cached
-from cachetools.keys import hashkey
+from cachetools import LRUCache, cached  # type: ignore[import-untyped]
+from cachetools.keys import hashkey  # type: ignore[import-untyped]
 from jsonschema.exceptions import ValidationError
 from pint.errors import DimensionalityError
 from pyisemail import is_email
@@ -47,7 +47,13 @@
 from tqdm.auto import tqdm
 from urllib3.util.retry import Retry
 
-from mp_api.client.contribs._types import Attachment, MPCDict, MPCStructure, Table
+from mp_api.client.contribs._types import (
+    Attachment,
+    MPCDict,
+    MPCStructure,
+    Table,
+    _Component,
+)
 from mp_api.client.contribs._units import ureg
 from mp_api.client.contribs.logger import MPCC_LOGGER, TqdmToLogger
 from mp_api.client.contribs.schemas import (
@@ -738,7 +744,7 @@ def query_projects(
         fields: list | None = None,
         sort: str | None = None,
         timeout: int = -1,
-    ) -> list[dict] | list[ContribsProject]:
+    ) -> list[MPCDict] | list[ContribsProject]:
         """Query projects by query and/or term (Atlas Search).
 
         See `client.available_query_params(resource="projects")` for keyword arguments used in
@@ -758,7 +764,7 @@ def query_projects(
         query = query or {}
 
         if self.project or "name" in query:
-            return self.get_project(name=query.get("name"), fields=fields)
+            return [self.get_project(name=query.get("name"), fields=fields)]  # type: ignore[return-value]
 
         if term:
 
@@ -1219,7 +1225,7 @@ def init_columns(
                                         f"Can't convert {existing_unit} to {new_unit} for {path}"
                                     )
                             try:
-                                factor = ureg.convert(1, *conv_args)
+                                factor = ureg.convert(1, *conv_args)  # type: ignore[arg-type]
                             except DimensionalityError:
                                 raise MPContribsClientError(
                                     f"Can't convert {existing_unit} to {new_unit} for {path}"
@@ -1704,7 +1710,7 @@ def query_contributions(
             ).result()
 
         if len(ret["data"]) > 0 and self.use_document_model:  # type: ignore[arg-type]
-            ret["data"] = [ContribData(**doc) for doc in ret["data"]]
+            ret["data"] = [ContribData(**doc) for doc in ret["data"]]  # type: ignore[arg-type,misc,union-attr]
 
         return (
             _convert_to_model(  # type: ignore[return-value]
@@ -2327,7 +2333,7 @@ def download_contributions(
 
                 match component:
                     case "structures":
-                        component_cls = MPCStructure
+                        component_cls: type[_Component] = MPCStructure
                     case "tables":
                         component_cls = Table
                     case "attachments":

mp_api/client/contribs/logger.py

@@ -94,4 +94,4 @@ def write(self, buf: str) -> int:
         return 1
 
     def flush(self) -> None:
-        self.logger.log(self.level, self.buf)
+        self.logger.log(self.level, self.buf)  # type: ignore[arg-type]

mp_api/client/contribs/schemas.py

@@ -3,7 +3,7 @@
 
 import re
 from datetime import datetime, timezone
-from typing import Any, Literal, Self, get_args
+from typing import TYPE_CHECKING, Any, Literal, Self, get_args
 
 import pandas as pd
 from pydantic import BaseModel, Field, create_model, field_serializer, field_validator
@@ -12,11 +12,14 @@
 from mp_api.client.contribs.utils import flatten_dict, unflatten_dict
 from mp_api.client.core.schemas import _DictLikeAccess
 
+if TYPE_CHECKING:
+    from types import UnionType
+
 CONTRIBS_DOC_NAME = "ContribsDoc"
 """Parent name of the dynamically-created contribs docs, similar to `MPDataDoc`."""
 
 
-def _cast_pandas_dtype(dtype: type, assume_nullable: bool = True) -> type:
+def _cast_pandas_dtype(dtype: type, assume_nullable: bool = True) -> type | UnionType:
     """Convert pandas dtype to built-in types.
 
     Args:
@@ -116,7 +119,7 @@ def _get_pydantic_from_dataframe(
         if not all(pd.isna(df[col_name]))
     }
 
-    return create_model("InferredModel", **model_fields), columns_renamed
+    return create_model("InferredModel", **model_fields), columns_renamed  # type: ignore[call-overload]
 
 
 class Reference(_DictLikeAccess):
@@ -239,7 +242,7 @@ def construct_data(cls, d: dict) -> dict[str, str | Datum]:
         }
         return {
             k: (
-                Datum(
+                Datum(  # type: ignore[misc]
                     **{
                         sub_k: flattened_data_dct.get(f"{k}.{sub_k}", field.default)
                         for sub_k, field in Datum.model_fields.items()

mp_api/client/mprester.py

@@ -247,7 +247,7 @@ def __init__(
                 )
 
     @property
-    def contribs(self) -> None:
+    def contribs(self):
         """Create an instance of the MP ContribsClient.
 
         NB: The following args are taken from `MPRester`: