feat: typechecking, candidate for 2.2.0

CHANGELOG.rst

@@ -1,12 +1,17 @@
 Changelog
 =========
 
+:Version: 2.2.0 of 2023-10-08
+Added type annotations (thanks to @charlesincharge for the initiative on an initial PR).
+Adapted type checking of size parameter for compatibility with mypy.  
+
 :Version: 2.1.0 of 2022-04-16
 Fix by @onnoeberhard for too-small dc-component: When cumulating the generated noise,
 the displacement would grow too slowly in the long limit.
 Test that would have discovered the above issue.
 Other tests are now deterministic.
 
+
 :Version: 2.0.0 of 2022-04-16
 
 Allow for control over random number generator state by adding optional random_state
@@ -18,9 +23,11 @@ Drop Python 2.7 support to use of NumPy's recommended default_rng constructor.
 Improve doc strings based on user questions.
 Check that fmin parameter is in the right range.
 
+
 :Version: 1.1.1 of 2019-02-08
 
-Use numpy's sum instead of python's (thank's to RuABraun).
+Use numpy's sum instead of python's (thanks to RuABraun).
+
 
 :Version: 1.1 of 2019-02-08
 

colorednoise.py

@@ -1,12 +1,18 @@
 """Generate colored noise."""
 
+from typing import Union, Iterable, Optional
 from numpy import sqrt, newaxis, integer
 from numpy.fft import irfft, rfftfreq
 from numpy.random import default_rng, Generator, RandomState
 from numpy import sum as npsum
 
 
-def powerlaw_psd_gaussian(exponent, size, fmin: float = 0.0, random_state=None):
+def powerlaw_psd_gaussian(
+        exponent: float, 
+        size: Union[int, Iterable[int]], 
+        fmin: float = 0.0, 
+        random_state: Optional[Union[int, Generator, RandomState]] = None
+    ):
     """Gaussian (1/f)**beta noise.
 
     Based on the algorithm in:
@@ -68,17 +74,19 @@ def powerlaw_psd_gaussian(exponent, size, fmin: float = 0.0, random_state=None):
     """
 
     # Make sure size is a list so we can iterate it and assign to it.
-    try:
-        size = list(size)
-    except TypeError:
+    if isinstance(size, (integer, int)):
         size = [size]
+    elif isinstance(size, Iterable):
+        size = list(size)
+    else:
+        raise ValueError("Size must be of type int or Iterable[int]")
 
     # The number of samples in each time series
     samples = size[-1]
 
     # Calculate Frequencies (we asume a sample rate of one)
     # Use fft functions for real output (-> hermitian spectrum)
-    f = rfftfreq(samples)
+    f = rfftfreq(samples) # type: ignore # mypy 1.5.1 has problems here 
 
     # Validate / normalise fmin
     if 0 <= fmin <= 0.5:
@@ -132,7 +140,7 @@ def powerlaw_psd_gaussian(exponent, size, fmin: float = 0.0, random_state=None):
     return y
 
 
-def _get_normal_distribution(random_state):
+def _get_normal_distribution(random_state: Optional[Union[int, Generator, RandomState]]):
     normal_dist = None
     if isinstance(random_state, (integer, int)) or random_state is None:
         random_state = default_rng(random_state)

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='colorednoise',
-    version='2.1.0',
+    version='2.2.0',
     description='Generate Gaussian (1/f)**beta noise (e.g. pink noise)',
     long_description="""Generate Gaussian distributed noise with a power law spectrum.
         Based on the algorithm in 
@@ -29,7 +29,7 @@
     keywords='1/f flicker power-law correlated colored noise generator',
     url='http://github.com/felixpatzelt/colorednoise',
     download_url=(
-      'https://github.com/felixpatzelt/colorednoise/archive/1.1.1.tar.gz'
+      'https://github.com/felixpatzelt/colorednoise/archive/2.2.0.tar.gz'
     ),
     author='Felix Patzelt',
     author_email='felix@neuro.uni-bremen.de',