"""Defines commin typings used throughout all submodules."""

import datetime
import os
import pathlib
import typing

import torch
import torch.nn
import torch.types

type Tensor = torch.Tensor
type Device = torch.device | torch.types.Device | typing.Literal["cpu", "cuda"]
type DType = torch.dtype


class TensorFactoryKWArgs(typing.TypedDict, total=False):
    """Keyword arguments for tensor factory functions."""

    device: Device
    dtype: DType
    pin_memory: bool


type StateDict = dict[str, Tensor]
type Size = torch.Size | tuple[int, ...]
type Buffer = bytes | bytearray | memoryview
type Pathable = str | pathlib.Path | os.PathLike
type Primitive = int | float | str | bytes | bytearray | memoryview
type Datetime = datetime.datetime
