Самый «Pythonic» способ заставить поведение, подобное массиву, на входных данных, не являющихся массивами?

Я работаю над личным проектом по вычислению и управлению целочисленными распределениями для применения в таких системах, как DnD и Pathfinder.

Я хочу создать функции, которые могут запрашивать функцию вероятностной массы (PMF) и кумулятивную функцию распределения (CDF) для определенных входных данных таким образом, чтобы они работали как для целочисленных, так и для входных данных, подобных массиву.

Я знаю, что для этой цели можно использовать N=np.asarray(N, dtype = int), и это, безусловно, самый чистый вариант, если он работает. Но внутренняя работа моей функции (в том виде, в котором она реализована, может быть, есть обходной путь...) требует, чтобы мой ввод был как минимум 1d с использованием N = np.atleast_1d(np.asarray(N, dtype = int)) (здесь важно приведение типов к int, отсюда и промежуточный шаг).

Теперь функция работает идеально, как и ожидалось, за исключением того случая, когда изначально ввод был просто int. Тип возвращаемого значения — это массив float с формой (1,), а не просто массив float 0d.

Мой обходной путь заключался в том, чтобы кэшировать полученную форму самой внутренней операции и изменить ее в конце. Шаблон моей функции может выглядеть так:

def my_func(N):
    N_shapecache = np.asarray(N, dtype = int)
    N = np.atleast_1d(N_shapecache)
    ret = do_stuff_that_requires_N_have_a_length(N)
    return ret.reshape(N_shapecache.shape)

Это идеально обеспечивает ожидаемое поведение для таких вещей, как операторы печати, где целочисленные входные данные при печати будут выглядеть как чистые числа с плавающей запятой (типом является массив 0d, но я думаю, что это хорошо)

Тем не менее, я думаю, что это также немного громоздко, и, возможно, есть лучший способ справиться с такими вещами в огромной библиотеке numpy, где после того, как я напишу свой собственный способ сделать что-то, в 90% случаев я понимаю, что есть встроенный numpy, которая делает то же самое в 40 раз быстрее. Но поиск не дал мне того, что я искал (обычно это одно из двух исправлений, которые я использую по отдельности, но не вместе).

Для полного контекста вот (в основном) минимальный рабочий блок кода, в котором есть моя функция.

import numpy as np

class integerDistribution:

    def __init__(self, min_val, probability_distribution, rtol = 1e-10):
        """Initializes the instance from a minimum value and a finite list of probabilities.

        Args:
            min_val : int
                The minimum value of our distribution. As an example, rolling a d20 
                would have a minimum value of 1, while the sum of 2d6 has a minimum 
                value of 2.
            probability_distribution : array-like of floats
                Each entry at index 'idx' in the array corresponds to the probability 
                of obtaining the value 'idx + min_val' from our distribution.
            rtol: float
                When initializing, we perform a sanity check that the sum of 
                probabilities is 1 to within a tolerance of rtol. Defaults to 1e-10

        Raises:
            ValueError: If probabilities do not sum to 1 within the specified
                tolerance (rtol)      
        """
        if not np.isclose(np.sum(probability_distribution), 1.0, rtol):
            raise ValueError(f"Probabilities summed to {np.sum(probability_distribution)}")

        self.min_val = int(min_val)
        self.max_val = self.min_val + len(probability_distribution) - 1
        self.values = self.min_val + np.arange(len(probability_distribution), dtype = int)
        self.probability_distribution = np.array(probability_distribution)
        

    
    def pmf(self, N):
        """Evaluates the probability mass function (PMF) at given value(s)

        The probability mass function is defined as p(N) = P(X = N), 
        where X is a random variable sampled from 'self', and N is the value 
        for which the probability is being calculated. 

        Args:
            N : int or list of int
                The value(s) for which to compute the probability. 
                Can be a single integer or a list/array of integers.

        Returns:
            p(N) : float or list of float
                The probability of obtaining each value in N when sampled from 'self'. 
                The return type matches the input type: a single float if N is an 
                integer, or a list/array of floats if N is a list/array of integers.
        """
        N_shapecache = np.asarray(N, dtype = int)
        N = np.atleast_1d(N_shapecache)
        p_N = np.zeros(N.shape)
        valid_indices = np.where((self.min_val <= N)*(N <= self.max_val))
        p_N[valid_indices] = self.probability_distribution[N[valid_indices] - self.min_val]
        return p_N.reshape(N_shapecache.shape)

🤔 А знаете ли вы, что...
Python является интерпретируемым языком программирования.