diff --git a/src/with_err/with_err.py b/src/with_err/with_err.py index fb49058..b7adfa5 100644 --- a/src/with_err/with_err.py +++ b/src/with_err/with_err.py @@ -1,9 +1,9 @@ # https://share.gemini.google/BDvazjX2RWsE +# https://chatgpt.com/share/6a8b20b2-4ad4-83ea-9926-4f204ebc4e65 +# https://chatgpt.com/share/6a8b209f-a6f4-83ea-a7da-5d4bf54e4ecd import inspect -import sys -import types -from collections.abc import Callable, Coroutine +from collections.abc import AsyncGenerator, Callable, Coroutine, Generator from functools import wraps from typing import Any, Protocol, overload @@ -17,39 +17,103 @@ def __call__(self, *args: P.args, **kwargs: P.kwargs) -> tuple[R | None, Excepti # @type_check_only class CoroutineWithErr[**P, R](Protocol): def __call__( - self, - *args: P.args, - **kwargs: P.kwargs, + self, *args: P.args, **kwargs: P.kwargs, ) -> Coroutine[Any, Any, tuple[R | None, Exception | None]]: ... +# @type_check_only +class GeneratorWithErr[**P, R](Protocol): + def __call__( + self, *args: P.args, **kwargs: P.kwargs, + ) -> Generator[tuple[R | None, Exception | None], None, None]: + ... + + +# @type_check_only +class AsyncGeneratorWithErr[**P, R](Protocol): + def __call__( + self, *args: P.args, **kwargs: P.kwargs + ) -> AsyncGenerator[tuple[R | None, Exception | None], None]: + ... + + +# @type_check_only +class Decorator(Protocol): + ''' + helper protocol for indirect decorators + XXX currently CallableWithErr is misclassfied as CoroutineWithErr + if Callable returns Any. + ''' + @overload + def __call__[**P, R]( + self, func: Callable[P, Coroutine[Any, Any, R]], / + ) -> CoroutineWithErr[P, R]: ... + + @overload + def __call__[**P, R]( + self, func: Callable[P, AsyncGenerator[R, Any]], / + ) -> AsyncGeneratorWithErr[P, R]: ... + + @overload + def __call__[**P, R]( + self, func: Callable[P, Generator[R, Any, Any]], / + ) -> GeneratorWithErr[P, R]: ... + + @overload + def __call__[**P, R]( + self, func: Callable[P, R], / + ) -> CallableWithErr[P, R]: + # XXX currently CallableWithErr is misclassfied as CoroutineWithErr + # if Callable returns Any. + ... + + @overload def with_err[**P, R]( - func: Callable[P, R], / -) -> CallableWithErr[P, R]: - # Overload 1: Called directly with a function -> with_err(func) + *exceptions: type[Exception], +) -> Decorator: + # Overload 1: called with exception types or no args -> with_err(*exceptions)(func) ... -# Async Direct @overload def with_err[**P, R]( - __func: Callable[P, Coroutine[Any, Any, R]], / + func: Callable[P, Coroutine[Any, Any, R]], / ) -> CoroutineWithErr[P, R]: - # Overload 2: Async call with a function -> with_err(func) + # Overload 4: Async call with a function -> with_err(func) + # XXX currently CallableWithErr is misclassfied as CoroutineWithErr + # if Callable returns Any. ... @overload def with_err[**P, R]( - *exceptions: type[Exception], -) -> Callable[[Callable[P, R]], CallableWithErr[P, R]]: - # Overload 3: Called with exception types or no args -> with_err(*exceptions)(func) + func: Callable[P, AsyncGenerator[R, Any]], / +) -> AsyncGeneratorWithErr[P, R]: + # Overload 2: async generator directly with a function -> with_err(func) ... -def with_err(*args): +@overload +def with_err[**P, R]( + func: Callable[P, Generator[R, Any, Any]], / +) -> GeneratorWithErr[P, R]: + # Overload 3: generator directly with a function -> with_err(func) + ... + + +@overload +def with_err[**P, R]( + func: Callable[P, R], / +) -> CallableWithErr[P, R]: + # Overload 5: Called directly with a function -> with_err(func) + # XXX currently CallableWithErr is misclassfied as CoroutineWithErr + # if Callable returns Any. + ... + + +def with_err[**P, R](*args): """ Wraps a function to return (result, Exception) instead of raising. """ @@ -71,49 +135,77 @@ def decorator(func): return decorator -def _make_wrapper[**P, R](func: Callable[P, R], exceptions: tuple[type[Exception], ...]): +def _make_wrapper(func, exceptions): if inspect.iscoroutinefunction(func): - @wraps(func) - async def async_wrapper(*args, **kwargs): + return _make_async_wrapper(func, exceptions) + elif inspect.isasyncgenfunction(func): + return _make_async_gen_wrapper(func, exceptions) + elif inspect.isgeneratorfunction(func): + return _make_sync_gen_wrapper(func, exceptions) + else: + return _make_sync_wrapper(func, exceptions) + + +def _make_sync_wrapper[**P, R]( + func: Callable[P, R], + exceptions: tuple[type[Exception], ...], +) -> CallableWithErr[P, R]: + @wraps(func) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> tuple[R | None, Exception | None]: + try: + return func(*args, **kwargs), None + except exceptions as e: + return None, e + return wrapper + + +def _make_async_wrapper[**P, R]( + func: Callable[P, Coroutine[Any, Any, R]], + exceptions: tuple[type[Exception], ...] +) -> CoroutineWithErr[P, R]: + @wraps(func) + async def async_wrapper(*args, **kwargs): + try: + result = await func(*args, **kwargs) + return result, None + except exceptions as e: + return None, e + return async_wrapper + + +def _make_sync_gen_wrapper[**P, R]( + func: Callable[P, Generator[R, Any, Any]], + exceptions: tuple[type[Exception], ...], +) -> GeneratorWithErr[P, R]: + @wraps(func) + def sync_gen_wrapper(*args, **kwargs): + gen = func(*args, **kwargs) + while True: try: - result = await func(*args, **kwargs) - return result, None + item = next(gen) + yield item, None + except StopIteration: + return except exceptions as err: - # 1. Fetch exception's original internal traceback. - tb = sys.exc_info()[2] - - # 2. Capture the caller frame executing func. - caller_frame = sys._getframe(1) - - # 3. Create a parent traceback frame and link it above 'tb'. - combined_tb = types.TracebackType( - tb_next=tb, - tb_frame=caller_frame, - tb_lasti=caller_frame.f_lasti, - tb_lineno=caller_frame.f_lineno, - ) - return None, err.with_traceback(combined_tb) - return async_wrapper - else: - @wraps(func) - def wrapper(*args: P.args, **kwargs: P.kwargs) -> tuple[R | None, Exception | None]: + yield None, err + return + return sync_gen_wrapper + + +def _make_async_gen_wrapper[**P, R]( + func: Callable[P, AsyncGenerator[R, Any]], + exceptions: tuple[type[Exception], ...], +) -> AsyncGeneratorWithErr[P, R]: + @wraps(func) + async def async_gen_wrapper(*args, **kwargs): + gen = func(*args, **kwargs) + while True: try: - return func(*args, **kwargs), None - except exceptions as e: - # 1. Fetch exception's original internal traceback. - tb = sys.exc_info()[2] - - # 2. Capture the caller frame executing func. - caller_frame = sys._getframe(1) - - # 3. Create a parent traceback frame and link it above 'tb'. - combined_tb = types.TracebackType( - tb_next=tb, - tb_frame=caller_frame, - tb_lasti=caller_frame.f_lasti, - tb_lineno=caller_frame.f_lineno - ) - - # 4. Attach the combined traceback back to the error instance - return None, e.with_traceback(combined_tb) - return wrapper + item = await anext(gen) + yield item, None + except StopAsyncIteration: + return + except exceptions as err: + yield None, err + return + return async_gen_wrapper diff --git a/tests/test_with_err.py b/tests/test_with_err.py index 498e352..e840d5c 100644 --- a/tests/test_with_err.py +++ b/tests/test_with_err.py @@ -1,3 +1,4 @@ +import inspect import json import re @@ -6,24 +7,59 @@ from with_err import get_err_strs, with_err -@with_err -def my_json_loads(a: str): - return json.loads(a) - - -def my_json_loads2(a: str): - return my_json_loads3(a) - - -def my_json_loads3(a: str): - return json.loads(a) - - def test_with_err_success(): ''' success. ''' json_loads_e = with_err()(json.loads) + signature = inspect.signature(json_loads_e) + sig_dict = { + name: { + 'obj': obj, + 'type': obj.kind, + 'default': obj.default, + 'annotation': obj.annotation, + 'name': obj.name, + } + for name, obj in signature.parameters.items()} + + print(f'sig_dict: {sig_dict}') + + # json.loads parameters + # very primitive info from inspect.signature. + # requiring typeshed_client to obtained python stdlib types. + assert 's' in sig_dict + assert sig_dict['s']['type'] == inspect._ParameterKind.POSITIONAL_OR_KEYWORD + assert sig_dict['s']['default'] is inspect._empty + assert sig_dict['s']['annotation'] is inspect._empty + assert 'cls' in sig_dict + assert sig_dict['cls']['type'] == inspect._ParameterKind.KEYWORD_ONLY + assert sig_dict['cls']['default'] is None + assert sig_dict['cls']['annotation'] is inspect._empty + assert 'object_hook' in sig_dict + assert sig_dict['object_hook']['type'] == inspect._ParameterKind.KEYWORD_ONLY + assert sig_dict['object_hook']['default'] is None + assert sig_dict['cls']['annotation'] is inspect._empty + assert 'parse_float' in sig_dict + assert sig_dict['parse_float']['type'] == inspect._ParameterKind.KEYWORD_ONLY + assert sig_dict['parse_float']['default'] is None + assert sig_dict['parse_float']['annotation'] is inspect._empty + assert 'parse_int' in sig_dict + assert sig_dict['parse_int']['type'] == inspect._ParameterKind.KEYWORD_ONLY + assert sig_dict['parse_int']['default'] is None + assert sig_dict['parse_int']['annotation'] is inspect._empty + assert 'parse_constant' in sig_dict + assert sig_dict['parse_constant']['type'] == inspect._ParameterKind.KEYWORD_ONLY + assert sig_dict['parse_constant']['default'] is None + assert sig_dict['parse_constant']['annotation'] is inspect._empty + assert 'object_pairs_hook' in sig_dict + assert sig_dict['object_pairs_hook']['type'] == inspect._ParameterKind.KEYWORD_ONLY + assert sig_dict['object_pairs_hook']['default'] is None + assert sig_dict['object_pairs_hook']['annotation'] is inspect._empty + assert 'kw' in sig_dict + assert sig_dict['kw']['type'] == inspect._ParameterKind.VAR_KEYWORD + assert sig_dict['kw']['default'] is inspect._empty + assert sig_dict['kw']['annotation'] is inspect._empty a = '{"test": 1}' the_struct, err = json_loads_e(a) @@ -46,10 +82,9 @@ def test_with_err_exception(): err_str = "\n".join(get_err_strs(err)) print(f'test_with_err: exception: err_str: {err_str}') - assert 'json.decoder.JSONDecodeError: Expecting value: line 1 column 10 (char 9)' in err_str - assert re.search(r'test_with_err.py", line \d+, in test_with_err_exception', err_str) assert re.search(r'with_err.py", line \d+, in wrapper', err_str) assert re.search(r'json/__init__.py", line \d+, in loads', err_str) + assert 'json.decoder.JSONDecodeError: Expecting value: line 1 column 10 (char 9)' in err_str def test_with_err_re_pattern_error_on_json(): @@ -63,6 +98,11 @@ def test_with_err_re_pattern_error_on_json(): json_loads_e(a) +@with_err +def my_json_loads(a: str): + return json.loads(a) + + def test_with_err_my_json_loads(): ''' decorator. @@ -75,10 +115,18 @@ def test_with_err_my_json_loads(): err_str = "\n".join(get_err_strs(err)) print(f'test_with_err: my_json_loads: err_str: {err_str}') - assert 'json.decoder.JSONDecodeError: Expecting value: line 1 column 10 (char 9)' in err_str - assert re.search(r'test_with_err.py", line \d+, in test_with_err_my_json_loads', err_str) assert re.search(r'with_err.py", line \d+, in wrapper', err_str) + assert re.search(r'test_with_err.py", line \d+, in my_json_loads', err_str) assert re.search(r'json/__init__.py", line \d+, in loads', err_str) + assert 'json.decoder.JSONDecodeError: Expecting value: line 1 column 10 (char 9)' in err_str + + +def my_json_loads2(a: str): + return my_json_loads3(a) + + +def my_json_loads3(a: str): + return json.loads(a) def test_with_err_my_json_loads2(): @@ -96,12 +144,11 @@ def test_with_err_my_json_loads2(): err_str = "\n".join(get_err_strs(err)) print(f'test_with_err: my_json_loads2: err_str: {err_str}') - assert 'json.decoder.JSONDecodeError: Expecting value: line 1 column 10 (char 9)' in err_str - assert re.search(r'test_with_err.py", line \d+, in test_with_err_my_json_loads2', err_str) assert re.search(r'with_err.py", line \d+, in wrapper', err_str) - assert re.search(r'test_with_err.py", line \d+, in my_json_loads3', err_str) assert re.search(r'test_with_err.py", line \d+, in my_json_loads2', err_str) + assert re.search(r'test_with_err.py", line \d+, in my_json_loads3', err_str) assert re.search(r'json/__init__.py", line \d+, in loads', err_str) + assert 'json.decoder.JSONDecodeError: Expecting value: line 1 column 10 (char 9)' in err_str def test_with_err_re_search(): @@ -159,9 +206,9 @@ async def test_with_err_async_err(): print(f'err_str: {err_str}') assert isinstance(err, ValueError) assert res is None - assert re.search(r'test_with_err.py", line \d+, in test_with_err_async_err', err_str) assert re.search(r'with_err.py", line \d+, in async_wrapper', err_str) assert re.search(r'test_with_err.py", line \d+, in async_fetch_data', err_str) + assert re.search(r'ValueError: Failed to reach endpoint', err_str) @pytest.mark.asyncio @@ -174,3 +221,118 @@ async def test_with_err_async_success(): assert err is None assert res == {'status': 'ok'} + + +@pytest.mark.asyncio +async def test_with_err_async_err2(): + ''' + test async err (ValueError) + ''' + async_fetch_data_e = with_err(ValueError)(async_fetch_data) + res, err = await async_fetch_data_e("bad") + err_stack = get_err_strs(err) + err_str = '\n'.join(err_stack) + print(f'err_str: {err_str}') + assert isinstance(err, ValueError) + assert res is None + assert re.search(r'with_err.py", line \d+, in async_wrapper', err_str) + assert re.search(r'test_with_err.py", line \d+, in async_fetch_data', err_str) + assert re.search(r'ValueError: Failed to reach endpoint', err_str) + + +@with_err +def my_stream(): + yield 1 + raise ValueError('invalid') + + +def test_with_err_yield(): + for idx, (each, err) in enumerate(my_stream()): + if idx == 0: + assert each == 1 + assert err is None + else: + assert each is None + assert isinstance(err, ValueError) + + +@with_err +def my_stream2(): + idx = 0 + + idx += 1 + yield idx + + idx += 1 + yield idx + + idx += 1 + yield idx + + +def test_with_err_yield_success(): + end_idx = None + for idx, (each, err) in enumerate(my_stream2()): + end_idx = idx + assert each == idx + 1 + assert err is None + + assert end_idx == 2 + + +@with_err +async def my_async_stream(): + yield 1 + raise ValueError('invalid') + + +@pytest.mark.asyncio +async def test_with_err_async_yield(): + async for each, err in my_async_stream(): + if each == 1: + assert each == 1 + assert err is None + else: + assert each is None + assert isinstance(err, ValueError) + + +@with_err +async def my_async_stream2(): + idx = 0 + + idx += 1 + yield idx + + idx += 1 + yield idx + + idx += 1 + yield idx + + +@pytest.mark.asyncio +async def test_with_err_async_yield_success(): + async for each, err in my_async_stream2(): + assert each in [1, 2, 3] + assert err is None + + +@pytest.mark.asyncio +async def test_with_err_async_yield_success2(): + gen = my_async_stream2() + + ret, err = await anext(gen) + assert ret == 1 + assert err is None + + ret, err = await anext(gen) + assert ret == 2 + assert err is None + + ret, err = await anext(gen) + assert ret == 3 + assert err is None + + with pytest.raises(StopAsyncIteration): + await anext(gen)