ytdl-sub/src/ytdl_sub/script/functions/array_functions.py

import itertools
from typing import List, Optional

from ytdl_sub.script.types.array import Array
from ytdl_sub.script.types.resolvable import (
    AnyArgument,
    Boolean,
    Integer,
    Lambda,
    LambdaReduce,
    LambdaTwo,
    Resolvable,
)
from ytdl_sub.script.utils.exceptions import (
    UNREACHABLE,
    ArrayValueDoesNotExist,
    FunctionRuntimeException,
)


class ArrayFunctions:
    @staticmethod
    def array(maybe_array: AnyArgument) -> Array:
        """
        :description:
          Tries to cast an unknown variable type to an Array.
        """
        if not isinstance(maybe_array, Array):
            raise FunctionRuntimeException(
                f"Tried and failed to cast {maybe_array.type_name()} as an Array"
            )
        return maybe_array

    @staticmethod
    def array_size(array: Array) -> Integer:
        """
        :description:
          Returns the size of an Array.
        """
        return Integer(len(array.value))

    @staticmethod
    def array_extend(*arrays: Array) -> Array:
        """
        :description:
          Combine multiple Arrays into a single Array.
        """
        output: List[Resolvable] = []
        for array in arrays:
            output.extend(array.value)

        return Array(output)

    @staticmethod
    def array_overlay(
        array: Array, overlap: Array, only_missing: Optional[Boolean] = None
    ) -> Array:
        """
        :description:
          Overlaps ``overlap`` onto ``array``. Can optionally only overlay missing indices.
        """
        output: List[Resolvable] = []
        output.extend(array.value)

        overlap_only_missing = only_missing and only_missing.value

        for idx, overlap_value in enumerate(overlap.value):
            if overlap_only_missing and idx < len(array.value):
                continue

            if idx < len(array.value):
                output[idx] = overlap_value
            else:
                output.append(overlap_value)

        return Array(output)

    @staticmethod
    def array_at(array: Array, idx: Integer, default: Optional[AnyArgument] = None) -> AnyArgument:
        """
        :description:
          Return the element in the Array at index ``idx``. If ``idx`` exceeds the array length,
          either return ``default`` if provided or throw an error.
        """
        try:
            return array.value[idx.value]
        except IndexError:
            if default is not None:
                return default
            raise

    @staticmethod
    def array_first(array: Array, fallback: AnyArgument) -> AnyArgument:
        """
        :description:
          Returns the first element whose boolean conversion is True. Returns fallback
          if all elements evaluate to False.
        """
        for val in array.value:
            if bool(val.value):
                return val

        return fallback

    @staticmethod
    def array_contains(array: Array, value: AnyArgument) -> Boolean:
        """
        :description:
          Return True if the value exists in the Array. False otherwise.
        """
        return Boolean(value in array.value)

    @staticmethod
    def array_index(array: Array, value: AnyArgument) -> Integer:
        """
        :description:
          Return the index of the value within the Array if it exists. If it does not, it will
          throw an error.
        """
        if not ArrayFunctions.array_contains(array=array, value=value):
            raise ArrayValueDoesNotExist(
                "Tried to get the index of a value in an Array that does not exist"
            )

        if isinstance(value, Resolvable):
            return Integer(array.value.index(value))

        raise UNREACHABLE

    @staticmethod
    def array_slice(array: Array, start: Integer, end: Optional[Integer] = None) -> Array:
        """
        :description:
          Returns the slice of the Array.
        """
        if end is not None:
            return Array(array.value[start.value : end.value])
        return Array(array.value[start.value :])

    @staticmethod
    def array_flatten(array: Array) -> Array:
        """
        :description:
          Flatten any nested Arrays into a single-dimensional Array.
        """
        output: List[Resolvable] = []
        for elem in array.value:
            if isinstance(elem, Array):
                output.extend(ArrayFunctions.array_flatten(elem).value)
            else:
                output.append(elem)

        return Array(output)

    @staticmethod
    def array_reverse(array: Array) -> Array:
        """
        :description:
          Reverse an Array.
        """
        return Array(list(reversed(array.value)))

    @staticmethod
    def array_product(*arrays: Array) -> Array:
        """
        :description:
          Returns the Cartesian product of elements from different arrays
        """
        out: List[Resolvable] = []
        for combo in itertools.product(*[arr.value for arr in arrays]):
            out.append(Array(combo))

        return Array(out)

    # pylint: disable=unused-argument

    @staticmethod
    def array_apply(array: Array, lambda_function: Lambda) -> Array:
        """
        :description:
          Apply a lambda function on every element in the Array.
        :usage:

        .. code-block:: python

           {
             %array_apply( [1, 2, 3] , %string )
           }

           # ["1", "2", "3"]
        """
        return Array([Array([val]) for val in array.value])

    @staticmethod
    def array_apply_fixed(
        array: Array,
        fixed_argument: AnyArgument,
        lambda2_function: LambdaTwo,
        reverse_args: Optional[Boolean] = None,
    ) -> Array:
        """
        :description:
          Apply a lambda function on every element in the Array, with ``fixed_argument``
          passed as a second argument to every invocation.
        """
        if reverse_args and reverse_args.value:
            return Array([Array([fixed_argument, val]) for val in array.value])

        return Array([Array([val, fixed_argument]) for val in array.value])

    @staticmethod
    def array_enumerate(array: Array, lambda_function: LambdaTwo) -> Array:
        """
        :description:
          Apply a lambda function on every element in the Array, where each arg
          passed to the lambda function is ``idx, element`` as two separate args.
        """
        return Array([Array([Integer(idx), val]) for idx, val in enumerate(array.value)])

    @staticmethod
    def array_reduce(array: Array, lambda_reduce_function: LambdaReduce) -> AnyArgument:
        """
        :description:
          Apply a reduce function on pairs of elements in the Array, until one element remains.
          Executes using the left-most and reduces in the right direction.
        """
        return array