Skip to content

Instantly share code, notes, and snippets.

@RandallPittmanOrSt
Last active July 12, 2024 20:24
Show Gist options
  • Save RandallPittmanOrSt/ce695f0b7d7717493649909d1926314a to your computer and use it in GitHub Desktop.
Save RandallPittmanOrSt/ce695f0b7d7717493649909d1926314a to your computer and use it in GitHub Desktop.
Proposal for annotating netCDF4.Variable
from typing import Any, Generic, Literal, Type, TypeVar, Union, overload
from typing_extensions import Self, TypeAlias
import numpy as np
import numpy.typing as npt
# --Just for this example--
import netCDF4
CompoundType = netCDF4.CompoundType
VLType = netCDF4.VLType
EnumType = netCDF4.EnumType
Dataset = netCDF4.Dataset
# --end--
# fmt: off
RealTypeLiteral: TypeAlias = Literal[
"i1", "b", "B", "int8", # NC_BYTE
"u1", "uint8", # NC_UBYTE
"i2", "h", "s", "int16", # NC_SHORT
"u2", "uint16", # NC_USHORT
"i4", "i", "l", "int32", # NC_INT
"u4", "uint32", # NC_UINT
"i8", "int64", "int", # NC_INT64
"u8", "uint64", # NC_UINT64
"f4", "f", "float32", # NC_FLOAT
"f8", "d", "float64", "float" # NC_DOUBLE
]
# fmt: on
ComplexTypeLiteral: TypeAlias = Literal["c8", "c16", "complex64", "complex128"]
NumericTypeLiteral: TypeAlias = RealTypeLiteral | ComplexTypeLiteral
CharTypeLiteral: TypeAlias = Literal["S1", "c"] # NC_CHAR
TypeLiteral: TypeAlias = NumericTypeLiteral | CharTypeLiteral
NPRealType: TypeAlias = np.int8 | np.uint8 | np.int16 | np.uint16 | np.int32 | np.uint32 | np.int64 | np.uint64 | np.float16 | np.float32 # nothing for S1
NPComplexType: TypeAlias = np.complex64 | np.complex128
NPNumType: TypeAlias = NPRealType | NPComplexType
VarDatatypeArgType: TypeAlias = TypeLiteral | Type[int | float | NPNumType | str | np.str_] | np.dtype | CompoundType | EnumType | VLType
VarT = TypeVar("VarT") # bound=Any
RealVarT = TypeVar("RealVarT", bound=NPRealType)
ComplexVarT = TypeVar("ComplexVarT", bound=NPComplexType)
NumericVarT = TypeVar("NumericVarT", bound=NPNumType)
# Dataset.createVar would be overloaded the same as Variable.__new__
class _VarDatatypeProperty:
# A descriptor definition of the property to allow overloads
@overload
def __get__(self, instance: Variable[RealVarT], owner: Any) -> RealVarT: ...
@overload
def __get__(self, instance: Variable[ComplexVarT], owner: Any) -> CompoundType: ...
@overload
def __get__(self, instance: Variable[str], owner: Any) -> VLType: ...
@overload
def __get__(self, instance: Variable[Any], owner: Any) -> Any: ... # actual return type np.dtype | CompoundType | VLType | EnumType
class _VarDtypeProperty:
# A descriptor definition of the property to allow overloads
@overload
def __get__(self, instance: Variable[NumericVarT], owner: Any) -> np.dtype[NumericVarT]: ...
@overload
def __get__(self, instance: Variable[str], owner: Any) -> Type[str]: ...
@overload
def __get__(self, instance: Variable[Any], owner: Any) -> Any: ... # actual return type np.dtype | Type[str]
class Variable(Generic[VarT]):
@overload
def __new__(cls, grp: Dataset, name: str, datatype: np.dtype[NumericVarT] | Type[NumericVarT], *args, **kwargs) -> Variable[NumericVarT]: ...
@overload
def __new__(cls, grp: Dataset, name: str, datatype: np.dtype[np.str_] | Type[str | np.str_], *args, **kwargs) -> Variable[str]: ...
@overload
def __new__(cls, grp: Dataset, name: str, datatype: VarDatatypeArgType, *args, **kwargs) -> Variable[Any]: ...
# Could also provide overloads for int-->np.int64, float-->np.float64. Also RealTypeLiteral returns some NPRealType
datatype: _VarDatatypeProperty
dtype: _VarDtypeProperty
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment