<commit_before>"""Array types extending numpy arrays.

The classes and types contained in here are primarily a way for PyType to reason
about. NumPy's typing system doesn't support immutability nor constrained
dimension sizes, so we introduce wrappers for common types (3D, 3D+Channels) to
allow for build-time error checking. This is generally preferable to seeing
errors at run-time during long-running pipelines.
"""

from collections import abc
import numbers
from typing import Any, Tuple, TypeVar, Type, Union

from connectomics.common import array_mixins
import numpy as np
import numpy.typing as npt

T = TypeVar('T', int, float)

# TODO(timblakely): Support Ellipses
ArbitrarySlice = Union[int, slice]
PointLookup = Union[np.ndarray, list[int]]
PointLookup4d = Tuple[ArbitrarySlice, PointLookup, PointLookup, PointLookup]
ArbitrarySlice4d = Tuple[ArbitrarySlice, ArbitrarySlice, ArbitrarySlice,
                         ArbitrarySlice]
CanonicalSlice4d = Tuple[slice, slice, slice, slice]

ArrayLike = Union[npt.ArrayLike, 'ImmutableArray', 'MutableArray']
Tuple3f = Tuple[float, float, float]
Tuple3i = Tuple[int, int, int]
Tuple4i = Tuple[int, int, int, int]
ArrayLike3d = Union[npt.ArrayLike, 'ImmutableArray', 'MutableArray', Tuple3f,
                    Tuple3i]


def normalize_index(ind: Union[ArbitrarySlice4d, PointLookup4d],
                    limits: Tuple4i) -> Union[CanonicalSlice4d, PointLookup4d]:
  """Converts a volume indexing expression into the canonical form.

  If the index expression is a point lookup, it is returned unchanged.

  Args:
    ind: indexing expression as passed to []
    limits: maximum limits in the event of open-ended index expression.

  Returns:
    4-sequence of slice objects for indexing subvolumes

  Raises:
    ValueError: when an unsupported indexing expression is passed
  """
  if not isinstance(ind, abc.Sized) or len(ind) != 4:
    raise ValueError('Only 4D slices are currently supported. '
                     f'Ellipses are unsupported. {ind} unsupported.')

  if all((isinstance(x, np.ndarray) or isinstance(x, list)) for x in ind[1:]):
    # Cast away the type here to make pytype happy. We've checked in the
    # conditional above that all the values for ind[1:] are either an ndarray or
    # a list.
    ind: Any = ind
    return ind

  ind = [process_slice_ind(ind_i, lim_i) for ind_i, lim_i in zip(ind, limits)]
  return tuple(ind)


def process_slice_ind(slice_ind: ArbitrarySlice, limit: int):
  """Converts a single slice index to canonical slice form.<commit_msg>Alias indexing expressions and point lookups.<commit_after>"""Array types extending numpy arrays.

The classes and types contained in here are primarily a way for PyType to reason
about. NumPy's typing system doesn't support immutability nor constrained
dimension sizes, so we introduce wrappers for common types (3D, 3D+Channels) to
allow for build-time error checking. This is generally preferable to seeing
errors at run-time during long-running pipelines.
"""

from collections import abc
import numbers
from typing import Any, Tuple, TypeVar, Type, Union

from connectomics.common import array_mixins
import numpy as np
import numpy.typing as npt

T = TypeVar('T', int, float)

# TODO(timblakely): Support Ellipses
ArbitrarySlice = Union[int, slice]
PointLookup = Union[np.ndarray, list[int]]
PointLookup4d = Tuple[ArbitrarySlice, PointLookup, PointLookup, PointLookup]
ArbitrarySlice4d = Tuple[ArbitrarySlice, ArbitrarySlice, ArbitrarySlice,
                         ArbitrarySlice]
CanonicalSlice4d = Tuple[slice, slice, slice, slice]
IndexExpOrPointLookup4d = Union[ArbitrarySlice4d, PointLookup4d]
CanonicalSliceOrPointLookup4d = Union[CanonicalSlice4d, PointLookup4d]

ArrayLike = Union[npt.ArrayLike, 'ImmutableArray', 'MutableArray']
Tuple3f = Tuple[float, float, float]
Tuple3i = Tuple[int, int, int]
Tuple4i = Tuple[int, int, int, int]
ArrayLike3d = Union[npt.ArrayLike, 'ImmutableArray', 'MutableArray', Tuple3f,
                    Tuple3i]


def normalize_index(ind: IndexExpOrPointLookup4d,
                    limits: Tuple4i) -> CanonicalSliceOrPointLookup4d:
  """Converts a volume indexing expression into the canonical form.

  If the index expression is a point lookup, it is returned unchanged.

  Args:
    ind: indexing expression as passed to []
    limits: maximum limits in the event of open-ended index expression.

  Returns:
    4-sequence of slice objects for indexing subvolumes

  Raises:
    ValueError: when an unsupported indexing expression is passed
  """
  if not isinstance(ind, abc.Sized) or len(ind) != 4:
    raise ValueError('Only 4D slices are currently supported. '
                     f'Ellipses are unsupported. {ind} unsupported.')

  if all((isinstance(x, np.ndarray) or isinstance(x, list)) for x in ind[1:]):
    # Cast away the type here to make pytype happy. We've checked in the
    # conditional above that all the values for ind[1:] are either an ndarray or
    # a list.
    ind: Any = ind
    return ind

  ind = [process_slice_ind(ind_i, lim_i) for ind_i, lim_i in zip(ind, limits)]
  return tuple(ind)


def process_slice_ind(slice_ind: ArbitrarySlice, limit: int):
  """Converts a single slice index to canonical slice form.
{'fraction_non_alphanumeric': 0.08040488922841864, 'fraction_numerical': 0.00878533231474408, 'mean_word_length': 4.20059582919563, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5741163', 'n_tokens_mistral': 1593, 'n_tokens_neox': 1511, 'n_words': 578}