Proposal for annotating netCDF4.Variable

VarTyping.pyi

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