revert hacky type check changes

Author: DanielYang59

pymatviz/bar.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, cast
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -11,24 +12,16 @@
 import plotly.graph_objects as go
 from matplotlib import transforms
 from matplotlib.ticker import FixedLocator
+from pymatgen.core import Structure
 from pymatgen.symmetry.groups import SpaceGroup
 
 from pymatviz.enums import Key
-from pymatviz.utils import (
-    PLOTLY,
-    Backend,
-    _check_type,
-    crystal_sys_from_spg_num,
-    si_fmt_int,
-)
+from pymatviz.utils import PLOTLY, Backend, crystal_sys_from_spg_num, si_fmt_int
 
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
     from typing import Any, Literal
 
-    from pymatgen.core import Structure
-
 
 def spacegroup_bar(
     data: Sequence[int | str | Structure] | pd.Series,
@@ -68,10 +61,10 @@ def spacegroup_bar(
     Returns:
         plt.Axes | go.Figure: matplotlib Axes or plotly Figure depending on backend.
     """
-    # TODO: use this hacky type check to avoid expensive import of Structure, #209
-    if _check_type(next(iter(data)), "pymatgen.core.structure.Structure"):
+    if isinstance(next(iter(data)), Structure):
         # if 1st sequence item is structure, assume all are
-        series = pd.Series(struct.get_space_group_info()[1] for struct in data)  # type: ignore[union-attr]
+        data = cast(Sequence[Structure], data)
+        series = pd.Series(struct.get_space_group_info()[1] for struct in data)
     else:
         series = pd.Series(data)
 

pymatviz/coordination.py

@@ -12,17 +12,17 @@
 import plotly.graph_objects as go
 from plotly.colors import label_rgb
 from plotly.subplots import make_subplots
+from pymatgen.analysis.local_env import NearNeighbors
 
 from pymatviz.colors import ELEM_COLORS_JMOL, ELEM_COLORS_VESTA
 from pymatviz.enums import ElemColorScheme, LabelEnum
-from pymatviz.utils import _check_type, normalize_to_dict
+from pymatviz.utils import normalize_to_dict
 
 
 if TYPE_CHECKING:
     from collections.abc import Callable
     from typing import Any, Literal
 
-    from pymatgen.analysis.local_env import NearNeighbors
     from pymatgen.core import PeriodicSite, Structure
 
 
@@ -66,13 +66,11 @@ def normalize_get_neighbors(
     if isinstance(strategy, int | float):
         return lambda site, structure: structure.get_neighbors(site, strategy)
 
-    if _check_type(strategy, "pymatgen.analysis.local_env.NearNeighbors"):
+    if isinstance(strategy, NearNeighbors):
         return lambda site, structure: strategy.get_nn_info(
             structure, structure.index(site)
         )
 
-    from pymatgen.analysis.local_env import NearNeighbors  # costly import
-
     if isclass(strategy) and issubclass(strategy, NearNeighbors):
         nn_instance = strategy()
         return lambda site, structure: nn_instance.get_nn_info(
@@ -430,8 +428,6 @@ def coordination_vs_cutoff_line(
     """
     structures = normalize_to_dict(structures)
 
-    from pymatgen.analysis.local_env import NearNeighbors
-
     # Determine cutoff range based on strategy
     if (
         isinstance(strategy, tuple)

pymatviz/phonons.py

@@ -10,6 +10,8 @@
 import plotly.graph_objects as go
 import scipy.constants as const
 from plotly.subplots import make_subplots
+from pymatgen.electronic_structure.bandstructure import BandStructureSymmLine
+from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine as PhononBands
 from pymatgen.util.string import htmlify
 
 
@@ -19,11 +21,10 @@
 
     import numpy as np
     from pymatgen.core import Structure
-    from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine as PhononBands
     from pymatgen.phonon.dos import PhononDos
     from typing_extensions import Self
 
-# AnyBandStructure: TypeAlias = BandStructureSymmLine | PhononBands
+AnyBandStructure: TypeAlias = BandStructureSymmLine | PhononBands
 YMin: TypeAlias = float | Literal["y_min"]
 YMax: TypeAlias = float | Literal["y_max"]
 BranchMode: TypeAlias = Literal["union", "intersection"]

pymatviz/powerups/both.py

@@ -8,6 +8,9 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import plotly.graph_objects as go
+import sklearn
+from sklearn.metrics import mean_absolute_percentage_error as mape
+from sklearn.metrics import r2_score
 
 from pymatviz.utils import (
     BACKENDS,
@@ -76,11 +79,6 @@ def annotate_metrics(
 
     backend: Backend = PLOTLY if isinstance(fig, go.Figure) else MATPLOTLIB
 
-    # Lazily import costly scikit-learn
-    import sklearn
-    from sklearn.metrics import mean_absolute_percentage_error as mape
-    from sklearn.metrics import r2_score
-
     funcs = {
         "MAE": lambda x, y: np.abs(x - y).mean(),
         "RMSE": lambda x, y: (((x - y) ** 2).mean()) ** 0.5,

pymatviz/process_data.py

@@ -7,6 +7,7 @@
 
 import pandas as pd
 from pandas.api.types import is_numeric_dtype, is_string_dtype
+from pymatgen.core import Composition
 
 from pymatviz.enums import ElemCountMode, Key
 from pymatviz.utils import ElemValues, df_ptable
@@ -62,8 +63,6 @@ def count_elements(
     # Ensure values is Series if we got dict/list/tuple
     srs = pd.Series(values)
 
-    from pymatgen.core import Composition  # costly import
-
     if is_numeric_dtype(srs):
         pass
 

pymatviz/ptable/ptable_plotly.py

@@ -8,6 +8,7 @@
 import numpy as np
 import pandas as pd
 import plotly.express as px
+import plotly.figure_factory as ff
 
 from pymatviz.enums import ElemCountMode
 from pymatviz.process_data import count_elements
@@ -271,8 +272,6 @@ def ptable_heatmap_plotly(
     zmax = max(non_nan_values) if cscale_range[1] is None else cscale_range[1]
     car_multiplier = 100 if heat_mode == "percent" else 1
 
-    import plotly.figure_factory as ff  # costly import
-
     fig = ff.create_annotated_heatmap(
         car_multiplier * heatmap_values,
         annotation_text=tile_texts,

pymatviz/relevance.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING
 
 import matplotlib.pyplot as plt
+import sklearn.metrics as skm
 
 from pymatviz.utils import df_to_arrays
 
@@ -36,8 +37,6 @@ def roc_curve(
     ax = ax or plt.gca()
 
     # get the metrics
-    import sklearn.metrics as skm
-
     false_pos_rate, true_pos_rate, _ = skm.roc_curve(targets, proba_pos)
     roc_auc = skm.roc_auc_score(targets, proba_pos)
 
@@ -72,8 +71,6 @@ def precision_recall_curve(
     ax = ax or plt.gca()
 
     # get the metrics
-    import sklearn.metrics as skm
-
     precision, recall, _ = skm.precision_recall_curve(targets, proba_pos)
 
     # proba_pos.round() converts class probabilities to integer class labels

pymatviz/utils.py

@@ -16,9 +16,11 @@
 import pandas as pd
 import plotly.graph_objects as go
 import plotly.io as pio
+import scipy.stats
 from matplotlib.colors import to_rgb
 from matplotlib.offsetbox import AnchoredText
 from matplotlib.ticker import FormatStrFormatter, PercentFormatter, ScalarFormatter
+from pymatgen.core import Structure
 
 
 if TYPE_CHECKING:
@@ -233,8 +235,6 @@ def bin_df_cols(
         )
 
     if density_col:
-        import scipy.stats  # expensive import
-
         # compute kernel density estimate for each bin
         values = df_in[bin_by_cols].dropna().T
         gaussian_kde = scipy.stats.gaussian_kde(values.astype(float))
@@ -679,7 +679,7 @@ def _get_matplotlib_font_color(fig: plt.Figure | plt.Axes) -> str:
 
 def normalize_to_dict(
     inputs: T | Sequence[T] | dict[str, T],
-    cls: type[T] | None = None,
+    cls: type[T] = Structure,
     key_gen: Callable[[T], str] = lambda obj: getattr(
         obj, "formula", type(obj).__name__
     ),
@@ -699,13 +699,8 @@ def normalize_to_dict(
     Raises:
         TypeError: If the input format is invalid.
     """
-    if cls is None:
-        from pymatgen.core import Structure  # costly import
-
-        cls = Structure
-
     if isinstance(inputs, cls):
-        return {"": inputs}  # type: ignore[dict-item]
+        return {"": inputs}
 
     if (
         isinstance(inputs, list | tuple)
@@ -728,23 +723,3 @@ def normalize_to_dict(
     raise TypeError(
         f"Invalid {inputs=}, expected {cls_name} or dict/list/tuple of {cls_name}"
     )
-
-
-def _check_type(obj: object, type_str: tuple[str, ...] | str) -> bool:
-    """Alternative to isinstance that avoids imports.
-
-    Todo:
-    Taken from monty.json, use until monty.json import fix merged.
-
-    Note for future developers: the type_str is not always obvious for an
-    object. For example, pandas.DataFrame is actually pandas.core.frame.DataFrame.
-    To find out the type_str for an object, run type(obj).mro(). This will
-    list all the types that an object can resolve to in order of generality
-    (all objects have the builtins.object as the last one).
-    """
-    type_str = type_str if isinstance(type_str, tuple) else (type_str,)
-    try:
-        mro = type(obj).mro()
-    except TypeError:
-        return False
-    return any(f"{o.__module__}.{o.__name__}" == ts for o in mro for ts in type_str)

pymatviz/xrd.py

@@ -8,17 +8,14 @@
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 from pymatgen.analysis.diffraction.xrd import DiffractionPattern, XRDCalculator
-
-from pymatviz.utils import _check_type
+from pymatgen.core import Structure
 
 
 if TYPE_CHECKING:
     from typing import Any, TypeAlias
 
-    from pymatgen.core import Structure
-
-    PatternOrStruct: TypeAlias = DiffractionPattern | Structure
 
+PatternOrStruct: TypeAlias = DiffractionPattern | Structure
 HklFormat: TypeAlias = Literal["compact", "full", None]
 ValidHklFormats = HklCompact, HklFull, HklNone = get_args(HklFormat)
 
@@ -101,9 +98,7 @@ def xrd_pattern(  # noqa: D417
         )
 
     # Convert single object to dict for uniform processing
-    if isinstance(patterns, DiffractionPattern) or _check_type(
-        patterns, "pymatgen.core.structure.Structure"
-    ):
+    if isinstance(patterns, DiffractionPattern | Structure):
         patterns = {"XRD Pattern": patterns}
     elif not isinstance(patterns, dict):
         raise TypeError(
@@ -141,7 +136,7 @@ def xrd_pattern(  # noqa: D417
         else:
             pattern_or_struct, trace_kwargs = pattern_data, {}
 
-        if _check_type(pattern_or_struct, "pymatgen.core.structure.Structure"):
+        if isinstance(pattern_or_struct, Structure):
             xrd_calculator = XRDCalculator(wavelength=wavelength)
             diffraction_pattern = xrd_calculator.get_pattern(pattern_or_struct)
         elif isinstance(pattern_or_struct, DiffractionPattern):