from __future__ import annotations
import abc
import enum
import re
from collections.abc import Iterable, Iterator, Sequence, Sized
from typing import NewType, Required, Self, TypeAlias, TypeVar, Union
import numpy
import pendulum
from numpy import (
bool_ as BoolDType,
)
from numpy import (
bytes_ as BytesDType,
)
from numpy import (
datetime64 as DateDType,
)
from numpy import (
float32 as FloatDType,
)
from numpy import (
generic as VarDType,
)
from numpy import (
int32 as IntDType,
)
from numpy import (
object_ as ObjDType,
)
from numpy import (
str_ as StrDType,
)
from numpy import (
uint8 as EnumDType,
)
from numpy.typing import DTypeLike, NDArray
from typing_extensions import Protocol, TypedDict
#: Generic covariant type var.
_T_co = TypeVar("_T_co", covariant=True)
# Arrays
#: Type var for numpy arrays.
_N_co = TypeVar("_N_co", covariant=True, bound="DTypeGeneric")
#: Type representing an numpy array.
Array: TypeAlias = NDArray[_N_co]
#: Type alias for a boolean array.
BoolArray: TypeAlias = Array[BoolDType]
#: Type alias for an array of bytes.
BytesArray: TypeAlias = Array[BytesDType]
#: Type alias for an array of dates.
DateArray: TypeAlias = Array[DateDType]
#: Type alias for an array of floats.
FloatArray: TypeAlias = Array[FloatDType]
#: Type alias for an array of enum indices.
IndexArray: TypeAlias = Array[EnumDType]
#: Type alias for an array of integers.
IntArray: TypeAlias = Array[IntDType]
#: Type alias for an array of objects.
ObjArray: TypeAlias = Array[ObjDType]
#: Type alias for an array of strings.
StrArray: TypeAlias = Array[StrDType]
#: Type alias for an array of generic objects.
VarArray: TypeAlias = Array[VarDType]
# Arrays-like
#: Type var for array-like objects.
_L = TypeVar("_L")
#: Type representing an array-like object.
ArrayLike: TypeAlias = Sequence[_L]
#: Type for bool arrays.
DTypeBool: TypeAlias = numpy.bool_
#: Type for int arrays.
DTypeInt: TypeAlias = numpy.int32
#: Type for float arrays.
DTypeFloat: TypeAlias = numpy.float32
#: Type for string arrays.
DTypeStr: TypeAlias = numpy.str_
#: Type for bytes arrays.
DTypeBytes: TypeAlias = numpy.bytes_
#: Type for Enum arrays.
DTypeEnum: TypeAlias = numpy.uint8
#: Type for date arrays.
DTypeDate: TypeAlias = numpy.datetime64
#: Type for "object" arrays.
DTypeObject: TypeAlias = numpy.object_
#: Type for "generic" arrays.
DTypeGeneric: TypeAlias = numpy.generic
# TODO(<Mauko Quiroga-Alvarado>): Properly resolve metaclass types.
# https://github.com/python/mypy/issues/14033
class _SeqIntMeta(type):
def __instancecheck__(self, arg: object, /) -> bool:
return (
bool(arg)
and isinstance(arg, Sequence)
and all(isinstance(item, int) for item in arg)
)
[docs]
class SeqInt(list[int], metaclass=_SeqIntMeta): ... # type: ignore[misc]
# Entities
#: For example "person".
EntityKey = NewType("EntityKey", str)
#: For example "persons".
EntityPlural = NewType("EntityPlural", str)
#: For example "principal".
RoleKey = NewType("RoleKey", str)
#: For example "parents".
RolePlural = NewType("RolePlural", str)
[docs]
class CoreEntity(Protocol):
key: EntityKey
plural: EntityPlural
def check_role_validity(self, role: object, /) -> None: ...
def check_variable_defined_for_entity(
self,
variable_name: VariableName,
/,
) -> None: ...
def get_variable(
self,
variable_name: VariableName,
check_existence: bool = ...,
/,
) -> None | Variable: ...
[docs]
class SingleEntity(CoreEntity, Protocol): ...
[docs]
class GroupEntity(CoreEntity, Protocol): ...
[docs]
class Role(Protocol):
entity: GroupEntity
max: int | None
subroles: None | Iterable[Role]
@property
def key(self, /) -> RoleKey: ...
@property
def plural(self, /) -> None | RolePlural: ...
# Indexed enums
[docs]
class EnumType(enum.EnumMeta):
indices: Array[DTypeEnum]
names: Array[DTypeStr]
enums: Array[DTypeObject]
[docs]
class Enum(enum.Enum, metaclass=EnumType):
index: int
_member_names_: list[str]
[docs]
class EnumArray(
numpy.ndarray[tuple[int, ...], numpy.dtype[DTypeEnum]], metaclass=abc.ABCMeta
):
possible_values: None | type[Enum]
@abc.abstractmethod
def __new__(
cls, input_array: Array[DTypeEnum], possible_values: type[Enum]
) -> Self: ...
# Holders
[docs]
class Holder(Protocol):
def clone(self, population: CorePopulation, /) -> Holder: ...
def get_memory_usage(self, /) -> MemoryUsage: ...
[docs]
class MemoryUsage(TypedDict, total=False):
cell_size: int
dtype: DTypeLike
nb_arrays: int
nb_cells_by_array: int
nb_requests: int
nb_requests_by_array: int
total_nb_bytes: Required[int]
# Parameters
#: A type representing a node of parameters.
ParameterNode: TypeAlias = Union[
"ParameterNodeAtInstant", "VectorialParameterNodeAtInstant"
]
#: A type representing a ???
ParameterNodeChild: TypeAlias = ParameterNode | ArrayLike[object]
[docs]
class ParameterNodeAtInstant(Protocol):
_instant_str: InstantStr
def __contains__(self, __item: object, /) -> bool: ...
def __getitem__(
self, __index: str | Array[DTypeGeneric], /
) -> ParameterNodeChild: ...
[docs]
class VectorialParameterNodeAtInstant(Protocol):
_instant_str: InstantStr
def __contains__(self, item: object, /) -> bool: ...
def __getitem__(
self, __index: str | Array[DTypeGeneric], /
) -> ParameterNodeChild: ...
# Periods
#: Matches "2015", "2015-01", "2015-01-01" but not "2015-13", "2015-12-32".
iso_format = re.compile(r"^\d{4}(-(?:0[1-9]|1[0-2])(-(?:0[1-9]|[12]\d|3[01]))?)?$")
#: Matches "2015", "2015-W01", "2015-W53-1" but not "2015-W54", "2015-W10-8".
iso_calendar = re.compile(r"^\d{4}(-W(0[1-9]|[1-4][0-9]|5[0-3]))?(-[1-7])?$")
#: For example 2020.
InstantInt = NewType("InstantInt", int)
#: For example 2020.
PeriodInt = NewType("PeriodInt", int)
class _InstantStrMeta(type):
def __instancecheck__(self, arg: object) -> bool:
return isinstance(arg, (ISOFormatStr, ISOCalendarStr))
[docs]
class InstantStr(str, metaclass=_InstantStrMeta): # type: ignore[misc]
__slots__ = ()
class _ISOFormatStrMeta(type):
def __instancecheck__(self, arg: object) -> bool:
return isinstance(arg, str) and bool(iso_format.match(arg))
class _ISOCalendarStrMeta(type):
def __instancecheck__(self, arg: object) -> bool:
return isinstance(arg, str) and bool(iso_calendar.match(arg))
[docs]
class ISOCalendarStr(str, metaclass=_ISOCalendarStrMeta): # type: ignore[misc]
__slots__ = ()
class _PeriodStrMeta(type):
def __instancecheck__(self, arg: object) -> bool:
return (
isinstance(arg, str)
and ":" in arg
and isinstance(arg.split(":")[1], InstantStr)
)
[docs]
class PeriodStr(str, metaclass=_PeriodStrMeta): # type: ignore[misc]
__slots__ = ()
[docs]
class Container(Protocol[_T_co]):
def __contains__(self, __item: object, /) -> bool: ...
[docs]
class Indexable(Protocol[_T_co]):
def __getitem__(self, __index: int, /) -> _T_co: ...
[docs]
class DateUnit(Container[str], Protocol):
def upper(self, /) -> str: ...
[docs]
class Instant(Indexable[int], Iterable[int], Sized, Protocol):
@property
def year(self, /) -> int: ...
@property
def month(self, /) -> int: ...
@property
def day(self, /) -> int: ...
@property
def date(self, /) -> pendulum.Date: ...
def __lt__(self, __other: object, /) -> bool: ...
def __le__(self, __other: object, /) -> bool: ...
def offset(self, __offset: str | int, __unit: DateUnit, /) -> None | Instant: ...
[docs]
class Period(Indexable[DateUnit | Instant | int], Protocol):
@property
def unit(self, /) -> DateUnit: ...
@property
def start(self, /) -> Instant: ...
@property
def size(self, /) -> int: ...
@property
def stop(self, /) -> Instant: ...
def contains(self, __other: Period, /) -> bool: ...
def offset(
self, __offset: str | int, __unit: None | DateUnit = None, /
) -> Period: ...
#: Type alias for a period-like object.
PeriodLike: TypeAlias = Period | PeriodStr | PeriodInt
# Populations
[docs]
class CorePopulation(Protocol): ...
[docs]
class SinglePopulation(CorePopulation, Protocol):
entity: SingleEntity
def get_holder(self, variable_name: VariableName, /) -> Holder: ...
[docs]
class GroupPopulation(CorePopulation, Protocol): ...
# Simulations
[docs]
class Simulation(Protocol):
def calculate(
self, variable_name: VariableName, period: Period, /
) -> Array[DTypeGeneric]: ...
def calculate_add(
self, variable_name: VariableName, period: Period, /
) -> Array[DTypeGeneric]: ...
def calculate_divide(
self, variable_name: VariableName, period: Period, /
) -> Array[DTypeGeneric]: ...
def get_population(self, plural: None | str, /) -> CorePopulation: ...
# Tax-Benefit systems
[docs]
class TaxBenefitSystem(Protocol):
person_entity: SingleEntity
def get_variable(
self,
variable_name: VariableName,
check_existence: bool = ...,
/,
) -> None | Variable: ...
# Tracers
#: A type representing a unit time.
Time: TypeAlias = float
#: A type representing a mapping of flat traces.
FlatNodeMap: TypeAlias = dict["NodeKey", "FlatTraceMap"]
#: A type representing a mapping of serialized traces.
SerializedNodeMap: TypeAlias = dict["NodeKey", "SerializedTraceMap"]
#: Key of a trace.
NodeKey = NewType("NodeKey", str)
[docs]
class FlatTraceMap(TypedDict, total=True):
dependencies: list[NodeKey]
parameters: dict[NodeKey, None | ArrayLike[object]]
value: None | VarArray
calculation_time: Time
formula_time: Time
[docs]
class SerializedTraceMap(TypedDict, total=True):
dependencies: list[NodeKey]
parameters: dict[NodeKey, None | ArrayLike[object]]
value: None | ArrayLike[object]
calculation_time: Time
formula_time: Time
[docs]
class SimpleTraceMap(TypedDict, total=True):
name: VariableName
period: int | Period
[docs]
class ComputationLog(Protocol):
def print_log(self, __aggregate: bool = ..., __max_depth: int = ..., /) -> None: ...
[docs]
class FlatTrace(Protocol):
def get_trace(self, /) -> FlatNodeMap: ...
def get_serialized_trace(self, /) -> SerializedNodeMap: ...
[docs]
class FullTracer(Protocol):
@property
def trees(self, /) -> list[TraceNode]: ...
def browse_trace(self, /) -> Iterator[TraceNode]: ...
def get_nb_requests(self, __name: VariableName, /) -> int: ...
[docs]
class SimpleTracer(Protocol):
@property
def stack(self, /) -> SimpleStack: ...
def record_calculation_start(
self, __name: VariableName, __period: PeriodInt | Period, /
) -> None: ...
def record_calculation_end(self, /) -> None: ...
[docs]
class TraceNode(Protocol):
@property
def children(self, /) -> list[TraceNode]: ...
@property
def end(self, /) -> Time: ...
@property
def name(self, /) -> str: ...
@property
def parameters(self, /) -> list[TraceNode]: ...
@property
def parent(self, /) -> None | TraceNode: ...
@property
def period(self, /) -> PeriodInt | Period: ...
@property
def start(self, /) -> Time: ...
@property
def value(self, /) -> None | VarArray: ...
def calculation_time(self, *, __round: bool = ...) -> Time: ...
def formula_time(self, /) -> Time: ...
def append_child(self, __node: TraceNode, /) -> None: ...
#: A stack of simple traces.
SimpleStack: TypeAlias = list[SimpleTraceMap]
# Variables
#: For example "salary".
VariableName = NewType("VariableName", str)
[docs]
class Variable(Protocol):
entity: CoreEntity
name: VariableName
[docs]
class Params(Protocol):
[docs]
def __call__(self, instant: Instant, /) -> ParameterNodeAtInstant: ...
__all__ = ["DTypeLike"]