fixes #787

Author: jph00

fastcore/_modidx.py

@@ -598,6 +598,7 @@
                                'fastcore.test.ExceptionExpected.__enter__': ('test.html#exceptionexpected.__enter__', 'fastcore/test.py'),
                                'fastcore.test.ExceptionExpected.__exit__': ('test.html#exceptionexpected.__exit__', 'fastcore/test.py'),
                                'fastcore.test.ExceptionExpected.__init__': ('test.html#exceptionexpected.__init__', 'fastcore/test.py'),
+                               'fastcore.test.expect_fail': ('test.html#expect_fail', 'fastcore/test.py'),
                                'fastcore.test.is_close': ('test.html#is_close', 'fastcore/test.py'),
                                'fastcore.test.nequals': ('test.html#nequals', 'fastcore/test.py'),
                                'fastcore.test.test': ('test.html#test', 'fastcore/test.py'),
@@ -748,6 +749,7 @@
                                 'fastcore.xtras.atomic_save': ('xtras.html#atomic_save', 'fastcore/xtras.py'),
                                 'fastcore.xtras.autostart': ('xtras.html#autostart', 'fastcore/xtras.py'),
                                 'fastcore.xtras.bunzip': ('xtras.html#bunzip', 'fastcore/xtras.py'),
+                                'fastcore.xtras.clean_cli_output': ('xtras.html#clean_cli_output', 'fastcore/xtras.py'),
                                 'fastcore.xtras.console_help': ('xtras.html#console_help', 'fastcore/xtras.py'),
                                 'fastcore.xtras.dataclass_src': ('xtras.html#dataclass_src', 'fastcore/xtras.py'),
                                 'fastcore.xtras.detect_mime': ('xtras.html#detect_mime', 'fastcore/xtras.py'),

fastcore/ansi.py

@@ -14,13 +14,14 @@
 
 def strip_terminal_queries(text):
     # Remove OSC sequences (like background color queries)
-    text = re.sub(r'\x1b\][^\\x07]*\x07', '', text)
+    text = re.sub('\x1b\\][^\x07]*\x07', '', text)
     # Remove DSR sequences (device status reports)
     return re.sub(r'\x1b\[[0-9]*n', '', text)
 
 
-def strip_ansi(source):
+def strip_ansi(source, term_queries:bool=False):
     "Remove ANSI escape codes from text."
+    if term_queries: source = strip_terminal_queries(source)
     return _ANSI_RE.sub("", source)
 
 

fastcore/basics.py

@@ -1098,11 +1098,12 @@ def __init__(self, f): self.f = f
     def __get__(self, _, f_cls): return MethodType(self.f, f_cls)
 
 # %% ../nbs/01_basics.ipynb #3f2733ef
-def patch_to(cls, as_prop=False, cls_method=False, set_prop=False, nm=None, glb=None):
+def patch_to(cls, as_prop=False, cls_method=False, set_prop=False, nm=None, glb=None, once:bool=False):
     "Decorator: add `f` to `cls`"
     if glb is None: glb = sys._getframe(1).f_globals
     def _inner(f):
         _nm = nm or f.__name__
+        onm = '_orig_'+_nm
         for c_ in tuplify(cls):
             nf = copy_func(f)
             for o in functools.WRAPPER_ASSIGNMENTS: setattr(nf, o, getattr(f,o))
@@ -1112,17 +1113,17 @@ def _inner(f):
             elif set_prop:    attr = getattr(c_, _nm).setter(nf)
             elif as_prop:     attr = property(nf)
             else:
-                onm = '_orig_'+_nm
-                if hasattr(c_, _nm) and not hasattr(c_, onm): setattr(c_, onm, getattr(c_, _nm))
+                if hasattr(c_, onm) and once: break
+                if hasattr(c_, _nm): setattr(c_, onm, getattr(c_, _nm))
                 attr = nf
             setattr(c_, _nm, attr)
         return glb.get(_nm, builtins.__dict__.get(_nm, None))
     return _inner
 
 # %% ../nbs/01_basics.ipynb #8faf7b86
-def patch(f=None, *, as_prop=False, cls_method=False, set_prop=False, nm=None):
+def patch(f=None, *, as_prop=False, cls_method=False, set_prop=False, nm=None, once:bool=False):
     "Decorator: add `f` to the first parameter's class (based on f's type annotations)"
-    if f is None: return partial(patch, as_prop=as_prop, cls_method=cls_method, set_prop=set_prop, nm=nm)
+    if f is None: return partial(patch, as_prop=as_prop, cls_method=cls_method, set_prop=set_prop, nm=nm, once=once)
     ann,glb,loc = get_annotations_ex(f)
     if cls_method:
         if 'cls' not in ann: raise TypeError(f"@patch with cls_method=True requires 'cls' to have a type annotation")
@@ -1131,7 +1132,7 @@ def patch(f=None, *, as_prop=False, cls_method=False, set_prop=False, nm=None):
         if not ann: raise TypeError(f"@patch requires the first parameter of `{f.__name__}` to have a type annotation")
         cls = next(iter(ann.values()))
     cls = union2tuple(eval_type(cls, glb, loc))
-    return patch_to(cls, as_prop=as_prop, cls_method=cls_method, set_prop=set_prop, nm=nm, glb=sys._getframe(1).f_globals)(f)
+    return patch_to(cls, as_prop=as_prop, cls_method=cls_method, set_prop=set_prop, nm=nm, once=once, glb=sys._getframe(1).f_globals)(f)
 
 # %% ../nbs/01_basics.ipynb #d1732261
 def compile_re(pat):

fastcore/test.py

@@ -3,14 +3,14 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_test.ipynb.
 
 # %% auto #0
-__all__ = ['TEST_IMAGE', 'TEST_IMAGE_BW', 'exception', 'test_fail', 'test', 'nequals', 'test_eq', 'test_eq_type', 'test_ne',
-           'is_close', 'test_close', 'test_is', 'test_shuffled', 'test_stdout', 'test_warns', 'test_fig_exists',
-           'ExceptionExpected']
+__all__ = ['TEST_IMAGE', 'TEST_IMAGE_BW', 'exception', 'test_fail', 'expect_fail', 'test', 'nequals', 'test_eq', 'test_eq_type',
+           'test_ne', 'is_close', 'test_close', 'test_is', 'test_shuffled', 'test_stdout', 'test_warns',
+           'test_fig_exists', 'ExceptionExpected']
 
 # %% ../nbs/00_test.ipynb #76d8df3a
 from .imports import *
 from collections import Counter
-from contextlib import redirect_stdout
+from contextlib import redirect_stdout,contextmanager
 
 # %% ../nbs/00_test.ipynb #2927d8a7
 def test_fail(f, msg='', contains='', exc=Exception, args=None, kwargs=None):
@@ -23,6 +23,16 @@ def test_fail(f, msg='', contains='', exc=Exception, args=None, kwargs=None):
     except Exception as e: assert False, f"Expected {exc.__name__} but got {type(e).__name__}: {e}. {msg}"
     assert False, f"Expected {exc.__name__} but none raised. {msg}"
 
+# %% ../nbs/00_test.ipynb #8197c1de
+@contextmanager
+def expect_fail(exc=Exception, contains=''):
+    "Context manager that fails unless body raises `exc` optionally containing `contains`"
+    try: yield
+    except exc as e:
+        assert not contains or contains in str(e)
+        return
+    assert False, f"Expected {exc.__name__} but none raised"
+
 # %% ../nbs/00_test.ipynb #a61150ba
 def test(a, b, cmp, cname=None):
     "`assert` that `cmp(a,b)`; display inputs and `cname or cmp.__name__` if it fails"

fastcore/xtras.py

@@ -10,19 +10,21 @@
            'detect_mime', 'bunzip', 'loads', 'loads_multi', 'dumps', 'untar_dir', 'repo_details', 'shell', 'ssh',
            'rsync_multi', 'run', 'open_file', 'save_pickle', 'load_pickle', 'parse_env', 'expand_wildcards',
            'atomic_save', 'dict2obj', 'obj2dict', 'repr_dict', 'is_listy', 'mapped', 'IterLen', 'ReindexCollection',
-           'SaveReturn', 'trim_wraps', 'save_iter', 'asave_iter', 'unqid', 'rtoken_hex', 'friendly_name',
-           'n_friendly_names', 'exec_eval', 'get_source_link', 'sparkline', 'modify_exception', 'round_multiple',
-           'set_num_threads', 'join_path_file', 'autostart', 'EventTimer', 'stringfmt_names', 'PartialFormatter',
-           'partial_format', 'truncstr', 'utc2local', 'local2utc', 'trace', 'modified_env', 'ContextManagers',
-           'shufflish', 'console_help', 'hl_md', 'type2str', 'dataclass_src', 'Unset', 'nullable_dc', 'make_nullable',
-           'flexiclass', 'asdict', 'vars_pub', 'is_typeddict', 'is_namedtuple', 'CachedIter', 'CachedAwaitable',
-           'reawaitable', 'is_async_callable', 'maybe_await', 'noopa', 'flexicache', 'time_policy', 'mtime_policy',
-           'timed_cache']
+           'SaveReturn', 'trim_wraps', 'save_iter', 'asave_iter', 'clean_cli_output', 'unqid', 'rtoken_hex',
+           'friendly_name', 'n_friendly_names', 'exec_eval', 'get_source_link', 'sparkline', 'modify_exception',
+           'round_multiple', 'set_num_threads', 'join_path_file', 'autostart', 'EventTimer', 'stringfmt_names',
+           'PartialFormatter', 'partial_format', 'truncstr', 'utc2local', 'local2utc', 'trace', 'modified_env',
+           'ContextManagers', 'shufflish', 'console_help', 'hl_md', 'type2str', 'dataclass_src', 'Unset', 'nullable_dc',
+           'make_nullable', 'flexiclass', 'asdict', 'vars_pub', 'is_typeddict', 'is_namedtuple', 'CachedIter',
+           'CachedAwaitable', 'reawaitable', 'is_async_callable', 'maybe_await', 'noopa', 'flexicache', 'time_policy',
+           'mtime_policy', 'timed_cache']
 
 # %% ../nbs/03_xtras.ipynb #3401d507
 from .imports import *
 from .foundation import *
 from .basics import *
+from .ansi import strip_ansi
+
 from importlib import import_module
 from functools import wraps
 import string,time,dataclasses
@@ -594,6 +596,16 @@ def asave_iter(g):
     def _(*args, **kwargs): return _save_iter(g, *args, **kwargs)
     return _
 
+# %% ../nbs/03_xtras.ipynb #45eb5141
+def clean_cli_output(txt:str, strip:bool=True):
+    "Clean CLI output by handling alternate screen, carriage returns, and ANSI escapes"
+    if '\x1b[?1049h' in txt:
+        if '\x1b[?1049l' not in txt: return ''
+        txt = txt.rsplit('\x1b[?1049l', 1)[-1]
+    txt = txt.replace('\r\n', '\n')
+    res = '\n'.join(l.rsplit('\r', 1)[-1] for l in txt.split('\n'))
+    return strip_ansi(res, term_queries=True) if strip else res
+
 # %% ../nbs/03_xtras.ipynb #35368de5
 def unqid(seeded=False):
     "Generate a unique id suitable for use as a Python identifier"

nbs/00_test.ipynb

@@ -20,7 +20,7 @@
     "#| export\n",
     "from fastcore.imports import *\n",
     "from collections import Counter\n",
-    "from contextlib import redirect_stdout"
+    "from contextlib import redirect_stdout,contextmanager"
    ]
   },
   {
@@ -116,12 +116,48 @@
    "outputs": [],
    "source": [
     "def _fail_args(a):\n",
-    "    if a == 5:\n",
-    "        raise ValueError\n",
+    "    if a == 5: raise ValueError\n",
     "test_fail(_fail_args, args=(5,))\n",
     "test_fail(_fail_args, kwargs=dict(a=5))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8197c1de",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "@contextmanager\n",
+    "def expect_fail(exc=Exception, contains=''):\n",
+    "    \"Context manager that fails unless body raises `exc` optionally containing `contains`\"\n",
+    "    try: yield\n",
+    "    except exc as e:\n",
+    "        assert not contains or contains in str(e)\n",
+    "        return\n",
+    "    assert False, f\"Expected {exc.__name__} but none raised\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "80939ade",
+   "metadata": {},
+   "source": [
+    "`expect_fail` is similar to `test_fail`, but is a context manager:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "49a1d0dd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with expect_fail(contains=\"foo\"): raise Exception(\"foobar\")\n",
+    "with expect_fail(ValueError): raise ValueError()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -851,13 +887,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "python3",
-   "language": "python",
-   "name": "python3"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 5
 }