Fixed several typing issues and added more thorough tests

Author: wolph

_python_utils_tests/test_aio.py

@@ -2,8 +2,9 @@
 import pytest
 import asyncio
 
+
 from python_utils import types
-from python_utils.aio import acount
+from python_utils.aio import acount, acontainer
 
 
 @pytest.mark.asyncio
@@ -20,3 +21,16 @@ async def mock_sleep(delay: float):
 
     assert len(sleeps) == 4
     assert sum(sleeps) == 4
+
+
+@pytest.mark.asyncio
+async def test_acontainer():
+    async def async_gen():
+        yield 1
+        yield 2
+        yield 3
+
+    assert await acontainer(async_gen) == [1, 2, 3]
+    assert await acontainer(async_gen()) == [1, 2, 3]
+    assert await acontainer(async_gen, set) == {1, 2, 3}
+    assert await acontainer(async_gen(), set) == {1, 2, 3}

_python_utils_tests/test_containers.py

@@ -3,8 +3,8 @@
 from python_utils import containers
 
 
-def test_unique_list_ignore():
-    a = containers.UniqueList()
+def test_unique_list_ignore() -> None:
+    a: containers.UniqueList[int] = containers.UniqueList()
     a.append(1)
     a.append(1)
     assert a == [1]
@@ -16,8 +16,10 @@ def test_unique_list_ignore():
     a[3] = 5
 
 
-def test_unique_list_raise():
-    a = containers.UniqueList(*range(20), on_duplicate='raise')
+def test_unique_list_raise() -> None:
+    a: containers.UniqueList[int] = containers.UniqueList(
+        *range(20), on_duplicate='raise'
+    )
     with pytest.raises(ValueError):
         a[10:20:2] = [1, 2, 3, 4, 5]
 

_python_utils_tests/test_decorators.py

@@ -2,37 +2,69 @@
 
 import pytest
 
-from python_utils.decorators import sample
+from python_utils.decorators import sample, wraps_classmethod
 
 
 @pytest.fixture
-def random(monkeypatch):
+def random(monkeypatch: pytest.MonkeyPatch) -> MagicMock:
     mock = MagicMock()
     monkeypatch.setattr(
-        "python_utils.decorators.random.random", mock, raising=True
+        'python_utils.decorators.random.random', mock, raising=True
     )
     return mock
 
 
-def test_sample_called(random):
+def test_sample_called(random: MagicMock):
     demo_function = MagicMock()
     decorated = sample(0.5)(demo_function)
     random.return_value = 0.4
     decorated()
     random.return_value = 0.0
     decorated()
     args = [1, 2]
-    kwargs = {"1": 1, "2": 2}
+    kwargs = {'1': 1, '2': 2}
     decorated(*args, **kwargs)
     demo_function.assert_called_with(*args, **kwargs)
     assert demo_function.call_count == 3
 
 
-def test_sample_not_called(random):
+def test_sample_not_called(random: MagicMock):
     demo_function = MagicMock()
     decorated = sample(0.5)(demo_function)
     random.return_value = 0.5
     decorated()
     random.return_value = 1.0
     decorated()
     assert demo_function.call_count == 0
+
+
+class SomeClass:
+    @classmethod
+    def some_classmethod(cls, arg):  # type: ignore
+        return arg  # type: ignore
+
+    @classmethod
+    def some_annotated_classmethod(cls, arg: int) -> int:
+        return arg
+
+
+def test_wraps_classmethod():  # type: ignore
+    some_class = SomeClass()
+    some_class.some_classmethod = MagicMock()
+    wrapped_method = wraps_classmethod(  # type: ignore
+        SomeClass.some_classmethod  # type: ignore
+    )(  # type: ignore
+        some_class.some_classmethod  # type: ignore
+    )
+    wrapped_method(123)
+    some_class.some_classmethod.assert_called_with(123)  # type: ignore
+
+
+def test_wraps_classmethod():  # type: ignore
+    some_class = SomeClass()
+    some_class.some_annotated_classmethod = MagicMock()
+    wrapped_method = wraps_classmethod(SomeClass.some_annotated_classmethod)(
+        some_class.some_annotated_classmethod
+    )
+    wrapped_method(123)  # type: ignore
+    some_class.some_annotated_classmethod.assert_called_with(123)

_python_utils_tests/test_generators.py

@@ -3,6 +3,7 @@
 import pytest
 
 import python_utils
+from python_utils import types
 
 
 @pytest.mark.asyncio
@@ -16,7 +17,7 @@ async def test_abatcher():
 
 @pytest.mark.asyncio
 async def test_abatcher_timed():
-    batches = []
+    batches: types.List[types.List[int]] = []
     async for batch in python_utils.abatcher(
         python_utils.acount(stop=10, delay=0.08), interval=0.1
     ):

_python_utils_tests/test_import.py

@@ -1,13 +1,13 @@
-from python_utils import import_
+from python_utils import import_, types
 
 
 def test_import_globals_relative_import():
     for i in range(-1, 5):
         relative_import(i)
 
 
-def relative_import(level):
-    locals_ = {}
+def relative_import(level: int):
+    locals_: types.Dict[str, types.Any] = {}
     globals_ = {'__name__': 'python_utils.import_'}
     import_.import_global('.formatters', locals_=locals_, globals_=globals_)
     assert 'camel_to_underscore' in globals_

_python_utils_tests/test_logger.py

@@ -15,5 +15,6 @@ class MyClass(Logurud):
     my_class.info('info')
     my_class.warning('warning')
     my_class.error('error')
+    my_class.critical('critical')
     my_class.exception('exception')
     my_class.log(0, 'log')

_python_utils_tests/test_time.py

@@ -5,6 +5,7 @@
 import pytest
 
 import python_utils
+from python_utils import types
 
 
 @pytest.mark.parametrize(
@@ -25,7 +26,12 @@
 )
 @pytest.mark.asyncio
 async def test_aio_timeout_generator(
-    timeout, interval, interval_multiplier, maximum_interval, iterable, result
+    timeout: float,
+    interval: float,
+    interval_multiplier: float,
+    maximum_interval: float,
+    iterable: types.AsyncIterable[types.Any],
+    result: int,
 ):
     i = None
     async for i in python_utils.aio_timeout_generator(
@@ -40,21 +46,30 @@ async def test_aio_timeout_generator(
     'timeout,interval,interval_multiplier,maximum_interval,iterable,result',
     [
         (0.01, 0.006, 0.5, 0.01, 'abc', 'c'),
-        (0.01, 0.006, 0.5, 0.01, itertools.count, 2),
+        (0.01, 0.006, 0.5, 0.01, itertools.count, 2),  # type: ignore
         (0.01, 0.006, 0.5, 0.01, itertools.count(), 2),
         (0.01, 0.006, 1.0, None, 'abc', 'c'),
         (
             timedelta(seconds=0.01),
             timedelta(seconds=0.006),
             2.0,
             timedelta(seconds=0.01),
-            itertools.count,
+            itertools.count,  # type: ignore
             2,
         ),
     ],
 )
 def test_timeout_generator(
-    timeout, interval, interval_multiplier, maximum_interval, iterable, result
+    timeout: float,
+    interval: float,
+    interval_multiplier: float,
+    maximum_interval: float,
+    iterable: types.Union[
+        str,
+        types.Iterable[types.Any],
+        types.Callable[..., types.Iterable[types.Any]],
+    ],
+    result: int,
 ):
     i = None
     for i in python_utils.timeout_generator(

pyproject.toml

@@ -4,8 +4,9 @@ target-version = ['py37', 'py38', 'py39', 'py310', 'py311']
 skip-string-normalization = true
 
 [tool.pyright]
-include = ['python_utils']
-strict = ['python_utils', '_python_utils_tests/test_aio.py']
+# include = ['python_utils']
+include = ['python_utils', '_python_utils_tests']
+strict = ['python_utils', '_python_utils_tests']
 # The terminal file is very OS specific and dependent on imports so we're skipping it from type checking
 ignore = ['python_utils/terminal.py']
-pythonVersion = '3.8'
+pythonVersion = '3.8'

python_utils/aio.py

@@ -8,6 +8,7 @@
 from . import types
 
 _N = types.TypeVar('_N', int, float)
+_T = types.TypeVar('_T')
 
 
 async def acount(
@@ -21,5 +22,33 @@ async def acount(
         if stop is not None and item >= stop:
             break
 
-        yield types.cast(_N, item)
+        yield item
         await asyncio.sleep(delay)
+
+
+async def acontainer(
+    iterable: types.Union[
+        types.AsyncIterable[_T],
+        types.Callable[..., types.AsyncIterable[_T]],
+    ],
+    container: types.Callable[[types.Iterable[_T]], types.Iterable[_T]] = list,
+) -> types.Iterable[_T]:
+    '''
+    Asyncio version of list()/set()/tuple()/etc() using an async for loop
+
+    So instead of doing `[item async for item in iterable]` you can do
+    `await acontainer(iterable)`.
+
+    '''
+    iterable_: types.AsyncIterable[_T]
+    if callable(iterable):
+        iterable_ = iterable()
+    else:
+        iterable_ = iterable
+
+    item: _T
+    items: types.List[_T] = []
+    async for item in iterable_:
+        items.append(item)
+
+    return container(items)

python_utils/decorators.py

@@ -3,9 +3,10 @@
 import random
 from . import types
 
-T = types.TypeVar('T')
-TC = types.TypeVar('TC', bound=types.Container[types.Any])
-P = types.ParamSpec('P')
+_T = types.TypeVar('_T')
+_TC = types.TypeVar('_TC', bound=types.Container[types.Any])
+_P = types.ParamSpec('_P')
+_S = types.TypeVar('_S', covariant=True)
 
 
 def set_attributes(**kwargs: types.Any) -> types.Callable[..., types.Any]:
@@ -33,8 +34,8 @@ def set_attributes(**kwargs: types.Any) -> types.Callable[..., types.Any]:
     '''
 
     def _set_attributes(
-        function: types.Callable[P, T]
-    ) -> types.Callable[P, T]:
+        function: types.Callable[_P, _T]
+    ) -> types.Callable[_P, _T]:
         for key, value in kwargs.items():
             setattr(function, key, value)
         return function
@@ -43,11 +44,13 @@ def _set_attributes(
 
 
 def listify(
-    collection: types.Callable[[types.Iterable[T]], TC] = list,  # type: ignore
+    collection: types.Callable[
+        [types.Iterable[_T]], _TC
+    ] = list,  # type: ignore
     allow_empty: bool = True,
 ) -> types.Callable[
-    [types.Callable[..., types.Optional[types.Iterable[T]]]],
-    types.Callable[..., TC],
+    [types.Callable[..., types.Optional[types.Iterable[_T]]]],
+    types.Callable[..., _TC],
 ]:
     '''
     Convert any generator to a list or other type of collection.
@@ -96,10 +99,10 @@ def listify(
     '''
 
     def _listify(
-        function: types.Callable[..., types.Optional[types.Iterable[T]]]
-    ) -> types.Callable[..., TC]:
-        def __listify(*args: types.Any, **kwargs: types.Any) -> TC:
-            result: types.Optional[types.Iterable[T]] = function(
+        function: types.Callable[..., types.Optional[types.Iterable[_T]]]
+    ) -> types.Callable[..., _TC]:
+        def __listify(*args: types.Any, **kwargs: types.Any) -> _TC:
+            result: types.Optional[types.Iterable[_T]] = function(
                 *args, **kwargs
             )
             if result is None:
@@ -134,10 +137,12 @@ def sample(sample_rate: float):
     '''
 
     def _sample(
-        function: types.Callable[P, T]
-    ) -> types.Callable[P, types.Optional[T]]:
+        function: types.Callable[_P, _T]
+    ) -> types.Callable[_P, types.Optional[_T]]:
         @functools.wraps(function)
-        def __sample(*args: P.args, **kwargs: P.kwargs) -> types.Optional[T]:
+        def __sample(
+            *args: _P.args, **kwargs: _P.kwargs
+        ) -> types.Optional[_T]:
             if random.random() < sample_rate:
                 return function(*args, **kwargs)
             else:
@@ -152,3 +157,43 @@ def __sample(*args: P.args, **kwargs: P.kwargs) -> types.Optional[T]:
         return __sample
 
     return _sample
+
+
+def wraps_classmethod(
+    wrapped: types.Callable[types.Concatenate[_S, _P], _T],
+) -> types.Callable[
+    [
+        types.Callable[types.Concatenate[types.Any, _P], _T],
+    ],
+    types.Callable[types.Concatenate[types.Type[_S], _P], _T],
+]:
+    '''
+    Like `functools.wraps`, but for wrapping classmethods with the type info
+    from a regular method
+    '''
+
+    def _wraps_classmethod(
+        wrapper: types.Callable[types.Concatenate[types.Any, _P], _T],
+    ) -> types.Callable[types.Concatenate[types.Type[_S], _P], _T]:
+        try:  # pragma: no cover
+            wrapper = functools.update_wrapper(
+                wrapper,
+                wrapped,
+                assigned=tuple(
+                    a
+                    for a in functools.WRAPPER_ASSIGNMENTS
+                    if a != '__annotations__'
+                ),
+            )
+        except AttributeError:
+            # For some reason `functools.update_wrapper` fails on some test
+            # runs but not while running actual code
+            pass
+
+        if annotations := getattr(wrapped, '__annotations__', {}):
+            annotations.pop('self', None)
+            wrapper.__annotations__ = annotations
+
+        return wrapper
+
+    return _wraps_classmethod