[Dev] Better type annotations

[DanielYang59]

As we're pushing towards more (static) type annotation coverage across the code base, I'm experiencing multiple issues and I would like to hear about your suggestions on the best practice here.

1. Use of "functionally similar" data types, for example as discussed here:
   For example the following Kpoint class annotated the kpts_shift arg as Vector3D(tuple[float, float, float]):

   pymatgen/pymatgen/io/vasp/inputs.py

   Lines 1024 to 1042 in 2500664

   	class Kpoints(MSONable):
   	"""KPOINT reader/writer."""
   	supported_modes = KpointsSupportedModes
   	def __init__(
   	self,
   	comment: str = "Default gamma",
   	num_kpts: int = 0,
   	style: KpointsSupportedModes = supported_modes.Gamma,
   	kpts: Sequence[float | Sequence] = ((1, 1, 1),),
   	kpts_shift: Vector3D = (0, 0, 0),
   	kpts_weights=None,
   	coord_type=None,
   	labels=None,
   	tet_number: int = 0,
   	tet_weight: float = 0,
   	tet_connections=None,
   	):

As a result codes that returns a list[float] would be flagged as having wrong types, for example multiple lines in the following that generate a list[float] instead of tuple[float, float, float]:

pymatgen/pymatgen/io/vasp/inputs.py

Lines 1371 to 1479 in 2500664

	@classmethod
	def from_str(cls, string):
	"""
	Reads a Kpoints object from a KPOINTS string.
	Args:
	string (str): KPOINTS string.
	Returns:
	Kpoints object
	"""
	lines = [line.strip() for line in string.splitlines()]
	comment = lines[0]
	num_kpts = int(lines[1].split()[0].strip())
	style = lines[2].lower()[0]
	# Fully automatic KPOINTS
	if style == "a":
	return cls.automatic(int(lines[3].split()[0].strip()))
	coord_pattern = re.compile(r"^\s*([\d+.\-Ee]+)\s+([\d+.\-Ee]+)\s+([\d+.\-Ee]+)")
	# Automatic gamma and Monk KPOINTS, with optional shift
	if style in ["g", "m"]:
	kpts = [int(i) for i in lines[3].split()]
	kpts_shift = (0, 0, 0)
	if len(lines) > 4 and coord_pattern.match(lines[4]):
	try:
	kpts_shift = [float(i) for i in lines[4].split()]
	except ValueError:
	pass
	return cls.gamma_automatic(kpts, kpts_shift) if style == "g" else cls.monkhorst_automatic(kpts, kpts_shift)
	# Automatic kpoints with basis
	if num_kpts <= 0:
	style = cls.supported_modes.Cartesian if style in "ck" else cls.supported_modes.Reciprocal
	kpts = [[float(j) for j in lines[i].split()] for i in range(3, 6)]
	kpts_shift = [float(i) for i in lines[6].split()]
	return Kpoints(
	comment=comment,
	num_kpts=num_kpts,
	style=style,
	kpts=kpts,
	kpts_shift=kpts_shift,
	)
	# Line-mode KPOINTS, usually used with band structures
	if style == "l":
	coord_type = "Cartesian" if lines[3].lower()[0] in "ck" else "Reciprocal"
	style = cls.supported_modes.Line_mode
	kpts = []
	labels = []
	patt = re.compile(r"([e0-9.\-]+)\s+([e0-9.\-]+)\s+([e0-9.\-]+)\s*!*\s*(.*)")
	for idx in range(4, len(lines)):
	line = lines[idx]
	m = patt.match(line)
	if m:
	kpts.append([float(m.group(1)), float(m.group(2)), float(m.group(3))])
	labels.append(m.group(4).strip())
	return Kpoints(
	comment=comment,
	num_kpts=num_kpts,
	style=style,
	kpts=kpts,
	coord_type=coord_type,
	labels=labels,
	)
	# Assume explicit KPOINTS if all else fails.
	style = cls.supported_modes.Cartesian if style in "ck" else cls.supported_modes.Reciprocal
	kpts = []
	kpts_weights = []
	labels = []
	tet_number = 0
	tet_weight = 0
	tet_connections = None
	for idx in range(3, 3 + num_kpts):
	tokens = lines[idx].split()
	kpts.append([float(j) for j in tokens[0:3]])
	kpts_weights.append(float(tokens[3]))
	if len(tokens) > 4:
	labels.append(tokens[4])
	else:
	labels.append(None)
	try:
	# Deal with tetrahedron method
	if lines[3 + num_kpts].strip().lower()[0] == "t":
	tokens = lines[4 + num_kpts].split()
	tet_number = int(tokens[0])
	tet_weight = float(tokens[1])
	tet_connections = []
	for idx in range(5 + num_kpts, 5 + num_kpts + tet_number):
	tokens = lines[idx].split()
	tet_connections.append((int(tokens[0]), [int(tokens[j]) for j in range(1, 5)]))
	except IndexError:
	pass
	return cls(
	comment=comment,
	num_kpts=num_kpts,
	style=cls.supported_modes[str(style)],
	kpts=kpts,
	kpts_weights=kpts_weights,
	tet_number=tet_number,
	tet_weight=tet_weight,
	tet_connections=tet_connections,
	labels=labels,

2. Handling of default values (especially None), for example in 3b293bf and

   pymatgen/pymatgen/io/vasp/optics.py

   Lines 139 to 193 in 2500664

   	def get_epsilon(
   	self,
   	idir: int,
   	jdir: int,
   	efermi: float | None = None,
   	nedos: int | None = None,
   	deltae: float | None = None,
   	ismear: int | None = None,
   	sigma: float | None = None,
   	cshift: float | None = None,
   	mask: NDArray | None = None,
   	) -> tuple[NDArray, NDArray]:
   	"""Compute the frequency dependent dielectric function.
   	Args:
   	idir: First direction of the dielectric tensor
   	jdir: Second direction of the dielectric tensor
   	efermi: Fermi energy
   	nedos: Number of points in the DOS
   	deltae: Energy step in the DOS
   	ismear: Smearing method (only has 0:gaussian, >0:Methfessel-Paxton)
   	sigma: Smearing width
   	cshift: Complex shift used for Kramer-Kronig transformation
   	mask: Mask for the bands/kpoint/spin index to include in the calculation
   	"""
   	def _use_default(param, default):
   	return param if param is not None else default
   	efermi = _use_default(efermi, self.efermi)
   	nedos = _use_default(nedos, self.nedos)
   	deltae = _use_default(deltae, self.deltae)
   	ismear = _use_default(ismear, self.ismear)
   	sigma = _use_default(sigma, self.sigma)
   	cshift = _use_default(cshift, self.cshift)
   	egrid, eps_imag = epsilon_imag( # type: ignore
   	cder=self.cder,
   	eigs=self.eigs,
   	kweights=self.kweights,
   	efermi=efermi, # type: ignore
   	nedos=nedos, # type: ignore
   	deltae=deltae, # type: ignore
   	ismear=ismear, # type: ignore
   	sigma=sigma, # type: ignore
   	idir=idir,
   	jdir=jdir,
   	mask=mask,
   	)
   	# scaling constant: edeps * np.pi / structure.volume
   	eps_in = eps_imag * edeps * np.pi / self.volume
   	eps = kramers_kronig(eps_in, nedos=nedos, deltae=deltae, cshift=cshift) # type: ignore
   	if idir == jdir:
   	eps += 1.0 + 0.0j
   	return egrid, eps

   . This would require explicit check for None, which would surely make the code more robust (but more "messy" as well).

Any advice would be appreciated. @janosh

[janosh]

sorry for the lack of reply, i completely missed this. are there any open questions left here?