from __future__ import annotations

__copyright__ = """

Copyright (C) 2009-2013 Andreas Kloeckner

Copyright (C) 2013- University of Illinois Board of Trustees

Copyright (C) 2020 Matt Wala

"""

__license__ = """

Permission is hereby granted, free of charge, to any person obtaining a copy

of this software and associated documentation files (the "Software"), to deal

in the Software without restriction, including without limitation the rights

to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

copies of the Software, and to permit persons to whom the Software is

furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in

all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

THE SOFTWARE.

"""

import builtins

import contextlib

import dataclasses

import logging

import operator

import re

import sys

from collections.abc import (

Callable,

Collection,

Hashable,

Iterable,

Iterator,

Mapping,

MutableSequence,

Sequence,

)

from functools import reduce, wraps

from pathlib import Path

from sys import intern

from typing import (

TYPE_CHECKING,

Any,

ClassVar,

Concatenate,

Generic,

Literal,

ParamSpec,

Protocol,

TypeVar,

cast,

overload,

)

from warnings import warn

from typing_extensions import (

Self,

TypeVarTuple,

Unpack,

dataclass_transform,

deprecated,

override,

)

from pytools.version import VERSION_TEXT

if TYPE_CHECKING:

import threading

import types

import numpy as np

from _typeshed import ReadableBuffer

from numpy.typing import DTypeLike, NDArray

from optype import CanLt

__version__ = VERSION_TEXT

# These are deprecated and will go away in 2022.

all = builtins.all

any = builtins.any

__doc__ = """

A Collection of Utilities

=========================

Math

----

.. autofunction:: levi_civita

Assertive accessors

-------------------

.. autofunction:: one

.. autofunction:: is_single_valued

.. autofunction:: all_roughly_equal

.. autofunction:: single_valued

Memoization

-----------

.. autofunction:: memoize

.. autofunction:: memoize_on_first_arg

.. autofunction:: memoize_method

.. autofunction:: memoize_in

.. autofunction:: keyed_memoize_on_first_arg

.. autofunction:: keyed_memoize_method

.. autofunction:: keyed_memoize_in

Argmin/max

----------

.. autofunction:: argmin2

.. autofunction:: argmax2

.. autofunction:: argmin

.. autofunction:: argmax

Reductions

----------

.. autofunction:: product

.. autofunction:: set_union

.. autofunction:: fset_union

Cartesian products

------------------

.. autofunction:: cartesian_product

.. autofunction:: distinct_pairs

Permutations, Tuples, Integer sequences

---------------------------------------

.. autofunction:: wandering_element

.. autofunction:: generate_nonnegative_integer_tuples_below

.. autofunction:: generate_nonnegative_integer_tuples_summing_to_at_most

.. autofunction:: generate_all_integer_tuples_below

.. autofunction:: generate_permutations

.. autofunction:: generate_unique_permutations

.. autoclass:: _ConcatenableSequence

Formatting

----------

.. autoclass:: Table

.. autofunction:: merge_tables

.. autofunction:: string_histogram

.. autofunction:: word_wrap

Debugging

---------

.. autofunction:: typedump

.. autofunction:: invoke_editor

Progress bars

-------------

.. autoclass:: ProgressBar

Name generation

---------------

.. autofunction:: generate_unique_names

.. autofunction:: generate_numbered_unique_names

.. autoclass:: UniqueNameGenerator

Deprecation Warnings

--------------------

.. autofunction:: module_getattr_for_deprecations

Functions for dealing with (large) auxiliary files

--------------------------------------------------

.. autofunction:: download_from_web_if_not_present

Helpers for :mod:`numpy`

------------------------

.. autofunction:: reshaped_view

Timing data

-----------

.. autoclass:: ProcessTimer

Log utilities

-------------

.. autoclass:: ProcessLogger

.. autoclass:: DebugProcessLogger

.. autoclass:: log_process

Sorting in natural order

------------------------

.. autofunction:: natorder

.. autofunction:: natsorted

Backports of newer Python functionality

---------------------------------------

.. autofunction:: resolve_name

Hashing

-------

.. autoclass:: Hash

.. autofunction:: unordered_hash

Sampling

--------

.. autofunction:: sphere_sample_equidistant

.. autofunction:: sphere_sample_fibonacci

String utilities

----------------

.. autofunction:: strtobool

.. autofunction:: to_identifier

Set-like functions for iterables

--------------------------------

These functions provide set-like operations on iterables. In contrast to

Python's built-in set type, they maintain the internal order of elements.

.. autofunction:: unique

.. autofunction:: unique_difference

.. autofunction:: unique_intersection

.. autofunction:: unique_union

Functionality for dataclasses

-----------------------------

.. autofunction:: opt_frozen_dataclass

Type Variables Used

-------------------

.. class:: T

Generic unbound invariant :class:`typing.TypeVar`.

.. class:: R_co

Generic unbound covariant :class:`typing.TypeVar`.

.. class:: F

Generic invariant :class:`typing.TypeVar` bound to a :class:`typing.Callable`.

.. class:: P

Generic unbound invariant :class:`typing.ParamSpec`.

.. class:: _HashT

.. class:: ReadableBuffer

Anything that implements the read-write buffer interface.

.. class:: SupportsLessThanT

An invariant :class:`typing.TypeVar` bound to :class:`optype.CanLt`.

"""

# {{{ type variables

T = TypeVar("T")

T_co = TypeVar("T_co", covariant=True)

T_contra = TypeVar("T_contra", contravariant=True)

Ts = TypeVarTuple("Ts")

R_co = TypeVar("R_co", covariant=True)

F = TypeVar("F", bound=Callable[..., Any])

P = ParamSpec("P")

K = TypeVar("K")

V = TypeVar("V")

EmptyT = TypeVar("EmptyT")

SupportsLessThanT = TypeVar("SupportsLessThanT", bound="CanLt[Any]")

# }}}

# {{{ code maintenance

# Undocumented on purpose for now, unclear that this is a great idea, given

# that typing.deprecated exists.

class MovedFunctionDeprecationWrapper(Generic[P, R_co]):

f: Callable[P, R_co]

deadline: int | str

def __init__(self,

f: Callable[P, R_co],

deadline: int | str | None = None) -> None:

if deadline is None:

deadline = "the future"

self.f = f

self.deadline = deadline

def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R_co:

from warnings import warn

warn(f"This function is deprecated and will go away in {self.deadline}. "

f"Use {self.f.__module__}.{self.f.__name__} instead.",

DeprecationWarning, stacklevel=2)

return self.f(*args, **kwargs)

# Deprecated. Should probably use `typing.deprecated` instead.

def deprecate_keyword(

oldkey: str,

newkey: str | None = None, *,

deadline: str | None = None

):

"""Decorator used to deprecate function keyword arguments.

:arg oldkey: deprecated argument name.

:arg newkey: new argument name that serves the same purpose, if any.

:arg deadline: expected time frame for the removal of the deprecated argument.

"""

from warnings import warn

if deadline is None:

deadline = "the future"

warn("'deprecate_keyword' is itself deprecated and will go away in Q1 2026. "

"Use 'warnings.deprecated' or a custom deprecation warning.",

DeprecationWarning, stacklevel=2)

def wrapper(func):

@wraps(func)

def inner_wrapper(*args, **kwargs):

if oldkey in kwargs:

if newkey is None:

warn(f"The '{oldkey}' keyword is deprecated and will "

f"go away in {deadline}.",

DeprecationWarning, stacklevel=2)

else:

warn(f"The '{oldkey}' keyword is deprecated and will "

f"go away in {deadline}. "

f"Use '{newkey}' instead.",

DeprecationWarning, stacklevel=2)

if newkey in kwargs:

raise ValueError(f"Cannot use '{oldkey}' "

f"and '{newkey}' in the same call.")

kwargs[newkey] = kwargs[oldkey]

del kwargs[oldkey]

return func(*args, **kwargs)

return inner_wrapper

return wrapper

def module_getattr_for_deprecations(

module_name: str,

depr_name_to_replacement_and_obj: Mapping[

str, tuple[str, object, str | int]

],

name: str

) -> object:

"""A helper to construct module-level :meth:`object.__getattr__` functions

so that deprecated names can still be found but raise a warning.

The typical usage pattern is as follows::

__getattr__ = partial(module_getattr_for_deprecations, __name__, {

"OldName": ("NewName", NewName, 2026),

})

"""

replacement_and_obj = depr_name_to_replacement_and_obj.get(name, None)

if replacement_and_obj is not None:

replacement, obj, deadline = replacement_and_obj

from warnings import warn

warn(f"'{module_name}.{name}' is deprecated. "

f"Use '{replacement}' instead. "

f"'{module_name}.{name}' will continue to work until {deadline}.",

DeprecationWarning, stacklevel=2)

return obj

raise AttributeError(name)

# }}}

# {{{ math

def delta(x: int, y: int) -> int:

if x == y:

return 1

return 0

def levi_civita(tup: tuple[int, ...]) -> int:

"""Compute an entry of the Levi-Civita symbol for the indices *tuple*."""

if len(tup) == 2:

i, j = tup

return j - i

if len(tup) == 3:

i, j, k = tup

return (j-i) * (k-i) * (k-j) // 2

raise NotImplementedError(f"Levi-Civita symbol in {len(tup)} dimensions")

def norm_1(iterable):

return sum(abs(x) for x in iterable)

def norm_2(iterable):

return sum(x**2 for x in iterable)**0.5

def norm_inf(iterable):

return max(abs(x) for x in iterable)

def norm_p(iterable, p):

return sum(i**p for i in iterable)**(1/p)

class Norm:

def __init__(self, p):

self.p = p

def __call__(self, iterable):

return sum(i**self.p for i in iterable)**(1/self.p)

# }}}

# {{{ data structures

# {{{ record

class RecordWithoutPickling:

"""An aggregate of named sub-variables. Assumes that each record sub-type

will be individually derived from this class.

"""

__slots__: ClassVar[list[str]] = []

fields: ClassVar[set[str]]

def __init__(self,

valuedict: Mapping[str, Any] | None = None,

exclude: Sequence[str] | None = None,

**kwargs: Any) -> None:

assert self.__class__ is not Record

if exclude is None:

exclude = ["self"]

try:

fields = self.__class__.fields

except AttributeError:

self.__class__.fields = fields = set()

if valuedict is not None:

kwargs.update(valuedict)

for key, value in kwargs.items():

if key not in exclude:

fields.add(key)

setattr(self, key, value)

def get_copy_kwargs(self, **kwargs):

for f in self.__class__.fields:

if f not in kwargs:

with contextlib.suppress(AttributeError):

kwargs[f] = getattr(self, f)

return kwargs

def copy(self, **kwargs) -> Self:

return self.__class__(**self.get_copy_kwargs(**kwargs))

@override

def __repr__(self):

return "{}({})".format(

self.__class__.__name__,

", ".join(f"{fld}={getattr(self, fld)!r}"

for fld in sorted(self.__class__.fields)

if hasattr(self, fld)))

def register_fields(self, new_fields):

try:

fields = self.__class__.fields

except AttributeError:

self.__class__.fields = fields = set()

fields.update(new_fields)

def __getattr__(self, name):

# This method is implemented to avoid pylint 'no-member' errors for

# attribute access.

raise AttributeError(

"'{}' object has no attribute '{}'".format(

self.__class__.__name__, name))

class Record(RecordWithoutPickling):

__slots__: ClassVar[list[str]] = []

def __getstate__(self):

return {

key: getattr(self, key)

for key in self.__class__.fields

if hasattr(self, key)}

def __setstate__(self, valuedict):

try:

fields = self.__class__.fields

except AttributeError:

self.__class__.fields = fields = set()

for key, value in valuedict.items():

fields.add(key)

setattr(self, key, value)

@override

def __eq__(self, other):

if self is other:

return True

return (self.__class__ == other.__class__

and self.__getstate__() == other.__getstate__())

@override

def __ne__(self, other):

return not self.__eq__(other)

class ImmutableRecordWithoutPickling(RecordWithoutPickling):

"""Hashable record. Does not explicitly enforce immutability."""

_cached_hash: int | None

def __init__(self, *args, **kwargs):

RecordWithoutPickling.__init__(self, *args, **kwargs)

self._cached_hash = None

@override

def __hash__(self) -> int:

# This attribute may vanish during pickling.

if getattr(self, "_cached_hash", None) is None:

self._cached_hash = hash((

type(self),

*(getattr(self, field) for field in self.__class__.fields)

))

assert self._cached_hash is not None

return self._cached_hash

class ImmutableRecord(ImmutableRecordWithoutPickling, Record):

pass

# }}}

class Reference(Generic[T]):

def __init__(self, value: T) -> None:

self.value: T = value

def get(self) -> T:

from warnings import warn

warn("Reference.get() is deprecated -- use ref.value instead. "

"This will stop working in 2025.", stacklevel=2)

return self.value

def set(self, value: T) -> None:

self.value = value

class FakeList(Generic[R_co]):

def __init__(self, f: Callable[[int], R_co], length: int) -> None:

self._Length: int = length

self._Function: Callable[[int], R_co] = f

def __len__(self) -> int:

return self._Length

@overload

def __getitem__(self, index: int) -> R_co: ...

@overload

def __getitem__(self, index: slice) -> Sequence[R_co]: ...

def __getitem__(self, index: object) -> R_co | Sequence[R_co]:

if isinstance(index, int):

return self._Function(index)

elif isinstance(index, slice):

return [self._Function(i) for i in range(*index.indices(self._Length))]

else:

raise TypeError(

f"list indices must be 'int' or 'slice', not '{type(index).__name__}'")

# {{{ dependent dictionary

class DependentDictionary(Generic[T, R_co]):

def __init__(self,

f: Callable[[dict[T, R_co], T], R_co],

start: dict[T, R_co] | None = None) -> None:

if start is None:

start = {}

self._Function: Callable[[dict[T, R_co], T], R_co] = f

self._Dictionary: dict[T, R_co] = start.copy()

def copy(self) -> DependentDictionary[T, R_co]:

return DependentDictionary(self._Function, self._Dictionary)

def __contains__(self, key: T) -> bool:

try:

self[key]

return True

except KeyError:

return False

def __getitem__(self, key: T) -> R_co:

try:

return self._Dictionary[key]

except KeyError:

return self._Function(self._Dictionary, key)

def __setitem__(self, key: T, value: R_co) -> None:

self._Dictionary[key] = value

def genuineKeys(self): # noqa: N802

return list(self._Dictionary.keys())

def iteritems(self) -> Iterable[tuple[T, R_co]]:

return self._Dictionary.items()

def iterkeys(self) -> Iterable[T]:

return self._Dictionary.keys()

def itervalues(self) -> Iterable[R_co]:

return self._Dictionary.values()

# }}}

# }}}

# {{{ assertive accessors

def one(iterable: Iterable[T]) -> T:

"""Return the first entry of *iterable*. Assert that *iterable* has only

that one entry.

"""

it = iter(iterable)

try:

v = next(it)

except StopIteration:

raise ValueError("empty iterable passed to 'one()'") from None

def no_more():

try:

next(it)

raise ValueError("iterable with more than one entry passed to 'one()'")

except StopIteration:

return True

assert no_more()

return v

def is_single_valued(

iterable: Iterable[T],

equality_pred: Callable[[T, T], bool] = operator.eq

) -> bool:

it = iter(iterable)

try:

first_item = next(it)

except StopIteration:

raise ValueError("empty iterable passed to 'single_valued()'") from None

return all(equality_pred(other_item, first_item) for other_item in it)

all_equal = is_single_valued

def all_roughly_equal(iterable: Iterable[object], threshold: float) -> bool:

return is_single_valued(iterable,

equality_pred=lambda a, b: abs(a-b) < threshold)

def single_valued(

iterable: Iterable[T],

equality_pred: Callable[[T, T], bool] = operator.eq

) -> T:

"""Return the first entry of *iterable*; Assert that other entries

are the same with the first entry of *iterable*.

"""

it = iter(iterable)

try:

first_item = next(it)

except StopIteration:

raise ValueError("empty iterable passed to 'single_valued()'") from None

assert all(equality_pred(other_item, first_item) for other_item in it)

return first_item

# }}}

# {{{ memoization / attribute storage

def memoize(*args: F, **kwargs: Any) -> F:

"""Stores previously computed function values in a cache.

Two keyword-only arguments are supported:

:arg use_kwargs: Allows the caller to use keyword arguments. Defaults to

``False``. Setting this to ``True`` has a non-negligible performance

impact.

:arg key: A function receiving the same arguments as the decorated function

which computes and returns the cache key.

"""

use_kw = bool(kwargs.pop("use_kwargs", False))

default_key_func: Callable[..., Any] | None

if use_kw:

def default_key_func(*inner_args, **inner_kwargs):

return inner_args, frozenset(inner_kwargs.items())

else:

default_key_func = None

key_func = kwargs.pop("key", default_key_func)

if kwargs:

raise TypeError(

"memoize received unexpected keyword arguments: {}".format(

", ".join(kwargs.keys())))

if key_func is not None:

def _decorator(func):

def wrapper(*args, **kwargs):

key = key_func(*args, **kwargs)

try:

return func._memoize_dic[key]

except AttributeError:

# _memoize_dic doesn't exist yet.

result = func(*args, **kwargs)

func._memoize_dic = {key: result}

return result

except KeyError:

result = func(*args, **kwargs)

func._memoize_dic[key] = result

return result

from functools import update_wrapper

update_wrapper(wrapper, func)

return wrapper

else:

def _decorator(func):

def wrapper(*args):

try:

return func._memoize_dic[args]

except AttributeError:

# _memoize_dic doesn't exist yet.

result = func(*args)

func._memoize_dic = {args: result}

return result

except KeyError:

result = func(*args)

func._memoize_dic[args] = result

return result

from functools import update_wrapper

update_wrapper(wrapper, func)

return wrapper

if not args:

return _decorator

if callable(args[0]) and len(args) == 1:

return _decorator(args[0])

raise TypeError(

f"memoize received unexpected position arguments: {args}")

FunctionValueCache = memoize

class _HasKwargs:

pass

def memoize_on_first_arg(

function: Callable[Concatenate[T, P], R_co], *,

cache_dict_name: str | None = None) -> Callable[Concatenate[T, P], R_co]:

"""Like :func:`memoize_method`, but for functions that take the object

in which do memoization information is stored as first argument.

Supports cache deletion via ``function_name.clear_cache(self)``.

"""

if cache_dict_name is None:

cache_dict_name = intern(

f"_memoize_dic_{function.__module__}{function.__name__}"

)

def wrapper(obj: T, *args: P.args, **kwargs: P.kwargs) -> R_co:

key = (_HasKwargs, frozenset(kwargs.items()), *args) if kwargs else args

assert cache_dict_name is not None

try:

return getattr(obj, cache_dict_name)[key]

except AttributeError:

attribute_error = True

except KeyError:

attribute_error = False

result = function(obj, *args, **kwargs)

if attribute_error:

object.__setattr__(obj, cache_dict_name, {key: result})

return result

getattr(obj, cache_dict_name)[key] = result

return result

def clear_cache(obj):

object.__delattr__(obj, cache_dict_name)

from functools import update_wrapper

new_wrapper = update_wrapper(wrapper, function)

# FIXME: pyright has a point here, stuffing random attributes into the

# function's dict is moderately sketchy. See what `functools.cache` is doing

# about it.

new_wrapper.clear_cache = clear_cache

return new_wrapper

def memoize_method(

method: Callable[Concatenate[T, P], R_co]

) -> Callable[Concatenate[T, P], R_co]:

"""Supports cache deletion via ``method_name.clear_cache(self)``.

.. versionchanged:: 2021.2

Can memoize methods on classes that do not allow setting attributes

(e.g. by overwriting ``__setattr__``), e.g. frozen :mod:`dataclasses`.

"""

return memoize_on_first_arg(method,

cache_dict_name=intern(f"_memoize_dic_{method.__name__}"))

class keyed_memoize_on_first_arg(Generic[T, P, R_co]): # noqa: N801

"""Like :func:`memoize_method`, but for functions that take the object

in which memoization information is stored as first argument.

Supports cache deletion via ``function_name.clear_cache(self)``.

:arg key: A function receiving the same arguments as the decorated function

which computes and returns the cache key.

:arg cache_dict_name: The name of the `dict` attribute in the instance

used to hold the cache.

.. versionadded :: 2020.3

"""

def __init__(self,

key: Callable[P, Hashable], *,

cache_dict_name: str | None = None) -> None:

self.key = key

self.cache_dict_name = cache_dict_name

def _default_cache_dict_name(self,

function: Callable[Concatenate[T, P], R_co]) -> str:

return intern(f"_memoize_dic_{function.__module__}{function.__name__}")

def __call__(

self, function: Callable[Concatenate[T, P], R_co]

) -> Callable[Concatenate[T, P], R_co]:

cache_dict_name = self.cache_dict_name

key = self.key

if cache_dict_name is None:

cache_dict_name = self._default_cache_dict_name(function)

def wrapper(obj: T, *args: P.args, **kwargs: P.kwargs) -> R_co:

cache_key = key(*args, **kwargs)

assert cache_dict_name is not None

try:

return getattr(obj, cache_dict_name)[cache_key]

except AttributeError:

result = function(obj, *args, **kwargs)

object.__setattr__(obj, cache_dict_name, {cache_key: result})

return result

except KeyError:

result = function(obj, *args, **kwargs)

getattr(obj, cache_dict_name)[cache_key] = result

return result

def clear_cache(obj):

object.__delattr__(obj, cache_dict_name)

from functools import update_wrapper

new_wrapper = update_wrapper(wrapper, function)

new_wrapper.clear_cache = clear_cache

return new_wrapper

class keyed_memoize_method(keyed_memoize_on_first_arg): # noqa: N801

"""Like :class:`memoize_method`, but additionally uses a function *key* to

compute the key under which the function result is stored.

Supports cache deletion via ``method_name.clear_cache(self)``.

:arg key: A function receiving the same arguments as the decorated function

which computes and returns the cache key.

.. versionadded :: 2020.3

.. versionchanged:: 2021.2

Can memoize methods on classes that do not allow setting attributes

(e.g. by overwriting ``__setattr__``), e.g. frozen :mod:`dataclasses`.

"""

@override

def _default_cache_dict_name(self, function):

return intern(f"_memoize_dic_{function.__name__}")

class memoize_in: # noqa: N801

"""Adds a cache to the function it decorates. The cache is attached

to *container* and must be uniquely specified by *identifier* (i.e.

all functions using the same *container* and *identifier* will be using

the same cache). The decorated function may only receive positional

arguments.

.. note::

This function works well on nested functions, which

do not have stable global identifiers.

.. versionchanged :: 2020.3

*identifier* no longer needs to be a :class:`str`,

but it needs to be hashable.

.. versionchanged:: 2021.2.1

Can now use instances of classes as *container* that do not allow

setting attributes (e.g. by overwriting ``__setattr__``),

e.g. frozen :mod:`dataclasses`.

"""

def __init__(self, container: Any, identifier: Hashable) -> None:

try:

memoize_in_dict = container._pytools_memoize_in_dict

except AttributeError:

memoize_in_dict = {}

object.__setattr__(container, "_pytools_memoize_in_dict",

memoize_in_dict)

self.cache_dict = memoize_in_dict.setdefault(identifier, {})

def __call__(self, inner: Callable[P, R_co]) -> Callable[P, R_co]:

@wraps(inner)

def new_inner(*args: P.args, **kwargs: P.kwargs) -> R_co:

assert not kwargs

try:

return self.cache_dict[args]

except KeyError:

result = inner(*args, **kwargs)

self.cache_dict[args] = result

return result

return new_inner