Skip to content

Commit

Permalink
Fix mypy errors on master branch (#3977)
Browse files Browse the repository at this point in the history
* improve type and simplify iterative_symmetrize

* fix type error and tweak types for compatibility

* fix packmol type

* fix type for site_transformation

* ignore type override

* fix type for species

* docstring tweak

* since it says preprocess in comment above, prob not meant to return here

---------

Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>
  • Loading branch information
DanielYang59 and janosh authored Aug 9, 2024
1 parent 0e90789 commit dbc68be
Show file tree
Hide file tree
Showing 8 changed files with 181 additions and 140 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ def __init__(self, lambda_table=None, alpha=-5):

# create Z and px
self.Z = 0
self._px: dict[Species, float] = defaultdict(float)
self._px: dict[SpeciesLike, float] = defaultdict(float)
for s1, s2 in itertools.product(self.species, repeat=2):
value = math.exp(self.get_lambda(s1, s2))
self._px[s1] += value / 2
Expand Down
20 changes: 15 additions & 5 deletions src/pymatgen/core/periodic_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from enum import Enum, unique
from itertools import combinations, product
from pathlib import Path
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, overload

import numpy as np
from monty.dev import deprecated
Expand Down Expand Up @@ -1597,11 +1597,21 @@ class DummySpecie(DummySpecies):
"""


@overload
def get_el_sp(obj: int) -> Element:
pass


@overload
def get_el_sp(obj: SpeciesLike) -> Element | Species | DummySpecies:
pass


@functools.lru_cache
def get_el_sp(obj: int | SpeciesLike) -> Element | Species | DummySpecies:
"""Utility method to get an Element, Species or DummySpecies from any input.
"""Utility function to get an Element, Species or DummySpecies from any input.
If obj is in itself an element or a specie, it is returned automatically.
If obj is an Element or a Species, it is returned as is.
If obj is an int or a string representing an integer, the Element with the
atomic number obj is returned.
If obj is a string, Species parsing will be attempted (e.g. Mn2+). Failing that
Expand All @@ -1616,8 +1626,8 @@ def get_el_sp(obj: int | SpeciesLike) -> Element | Species | DummySpecies:
ValueError: if obj cannot be converted into an Element or Species.
Returns:
Species | Element: with a bias for the maximum number of properties
that can be determined.
Element | Species | DummySpecies: with a bias for the maximum number
of properties that can be determined.
"""
# If obj is already an Element or Species, return as is
if isinstance(obj, (Element, Species, DummySpecies)):
Expand Down
Loading

0 comments on commit dbc68be

Please sign in to comment.