diff --git a/.github/workflows/pytest_and_autopublish.yml b/.github/workflows/pytest_and_autopublish.yml index 65a0b5c9..18dba91a 100644 --- a/.github/workflows/pytest_and_autopublish.yml +++ b/.github/workflows/pytest_and_autopublish.yml @@ -22,7 +22,7 @@ jobs: # Install deps - uses: actions/setup-python@v4 with: - python-version: 3.9 + python-version: "3.10" - run: pip --version - run: pip install -e .[all,dev] - run: pip freeze diff --git a/CHANGELOG.md b/CHANGELOG.md index 7e4c1910..a04fe9be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,9 @@ Changelog follow https://keepachangelog.com/ format. * `ecolab`: * Added protobuf repeated fields support to `ecolab.inspect` - * Fixed auto-display when the line contain UTF-8 character + * `ecolab.auto_display`: + * Support specifiers to customize auto-display (`;s`, `;a`, `;i`,...) + * Fixed auto-display when the line contain UTF-8 character ## [1.5.2] - 2023-10-24 diff --git a/etils/ecolab/README.md b/etils/ecolab/README.md index 1d449ac7..f7bf5d71 100644 --- a/etils/ecolab/README.md +++ b/etils/ecolab/README.md @@ -73,14 +73,25 @@ Add a trailing `;` to any statement (assignment, expression, return statement) to display the current line. This call `IPython.display.display()` for pretty display. +Format: + +* `my_obj;`: Alias for `IPython.display.display(x)` +* `my_obj;s`: (`spec`) Alias for `IPython.display.display(etree.spec_like(x))` +* `my_obj;i`: (`inspect`) Alias for `ecolab.inspect(x)` +* `my_obj;a`: (`array`) Alias for `media.show_images(x)` / + `media.show_videos(x)` (`ecolab.auto_plot_array` behavior) +* `my_obj;q`: (`quiet`) Don't display the line (e.g. last line) + ```python x = my_fn(); # Display `my_fn()` output my_fn(); # Display `my_fn()` output +my_fn();i # Inspect `my_fn()` output ``` -Note that `;` added to the last statement of the cell still silence the -output (`IPython` default behavior). +Note that contrary to `IPython` default behavior, `;` added to the last +statement of the cell will display the line. To silence the last output, use +`;q`. `;` behavior can be disabled with `ecolab.auto_display(False)` diff --git a/etils/ecolab/array_as_img.py b/etils/ecolab/array_as_img.py index 13cdd306..959c4544 100644 --- a/etils/ecolab/array_as_img.py +++ b/etils/ecolab/array_as_img.py @@ -53,6 +53,7 @@ def show(*objs, **kwargs) -> None: def auto_plot_array( *, + # If updating this, also update `_array_repr_html_inner` !!! video_min_num_frames: int = 15, # Images outside this range are rescalled height: None | int | tuple[int, int] = (100, 250), @@ -79,11 +80,8 @@ def auto_plot_array( if ipython is None: return # Non-notebook environement - show_images_kwargs = show_images_kwargs or {} - show_videos_kwargs = show_videos_kwargs or {} - array_repr_html_fn = functools.partial( - _array_repr_html, + array_repr_html, video_min_num_frames=video_min_num_frames, height=height, show_images_kwargs=show_images_kwargs, @@ -126,7 +124,7 @@ def auto_plot_array( formatter.for_type(enp.lazy.np.ndarray, array_repr_html_fn) -def _array_repr_html( +def array_repr_html( array: Array, **kwargs: Any, ) -> Optional[str]: @@ -142,12 +140,15 @@ def _array_repr_html( def _array_repr_html_inner( img: Array, *, - video_min_num_frames: int, - height: None | int | tuple[int, int], - show_images_kwargs: dict[str, Any], - show_videos_kwargs: dict[str, Any], + # If updating this, also update `auto_plot_array` !!! + video_min_num_frames: int = 15, + height: None | int | tuple[int, int] = (100, 250), + show_images_kwargs: Optional[dict[str, Any]] = None, + show_videos_kwargs: Optional[dict[str, Any]] = None, ) -> Optional[str]: """Display the normalized img, or `None` if the input is not an image.""" + show_images_kwargs = show_images_kwargs or {} + show_videos_kwargs = show_videos_kwargs or {} if not enp.lazy.is_array(img): # Not an array return None diff --git a/etils/ecolab/array_as_img_test.py b/etils/ecolab/array_as_img_test.py index ef54ebf5..f10a434d 100644 --- a/etils/ecolab/array_as_img_test.py +++ b/etils/ecolab/array_as_img_test.py @@ -57,7 +57,7 @@ def test_array_repr_html_valid( valid_shape: tuple[int, ...], ): # 2D images are displayed as images - assert ' None: statement) to display the current line. This call `IPython.display.display()` for pretty display. - Note that `;` added to the last statement of the cell still silence the - output (`IPython` default behavior). + This change the default IPython behavior where `;` added to the last statement + of the cell still silence the output. ```python x = my_fn(); # Display `my_fn()` @@ -51,6 +57,16 @@ def auto_display(activate: bool = True) -> None: `;` behavior can be disabled with `ecolab.auto_display(False)` + Format: + + * `my_obj;`: Alias for `IPython.display.display(x)` + * `my_obj;s`: (`spec`) Alias for + `IPython.display.display(etree.spec_like(x))` + * `my_obj;i`: (`inspect`) Alias for `ecolab.inspect(x)` + * `my_obj;a`: (`array`) Alias for `media.show_images(x)` / + `media.show_videos(x)` (`ecolab.auto_plot_array` behavior) + * `my_obj;q`: (`quiet`) Don't display the line (e.g. last line) + Args: activate: Allow to disable `auto_display` """ @@ -116,6 +132,8 @@ class _RecordLines(IPython.core.inputtransformer.InputTransformer): def __init__(self): self._lines = [] self.last_lines = [] + + self.trailing_stmt_line_nums = {} super().__init__() def push(self, line): @@ -132,6 +150,8 @@ def reset(self): for line in self._lines: self.last_lines.extend(line.split('\n')) self._lines.clear() + + self.trailing_stmt_line_nums.clear() return else: @@ -141,9 +161,13 @@ class _RecordLines: def __init__(self): self.last_lines = [] + # Additional state (reset at each cell) to keep track of which lines + # contain trailing statements + self.trailing_stmt_line_nums = {} def __call__(self, lines: list[str]) -> list[str]: self.last_lines = [l.rstrip('\n') for l in lines] + self.trailing_stmt_line_nums = {} return lines @@ -159,22 +183,25 @@ def _maybe_display( ) -> ast.AST: """Wrap the node in a `display()` call.""" try: - has_trailing, is_last_statement = _has_trailing_semicolon( - self.lines_recorder.last_lines, node - ) - if has_trailing: - if is_last_statement and isinstance(node, ast.Expr): - # Last expressions are already displayed by IPython, so instead - # IPython silence the statement - pass - elif node.value is None: # `AnnAssign().value` can be `None` (`a: int`) + if self._is_alias_stmt(node): # Alias statements should be no-op + return ast.Pass() + + line_info = _has_trailing_semicolon(self.lines_recorder.last_lines, node) + if line_info.has_trailing: + if node.value is None: # `AnnAssign().value` can be `None` (`a: int`) pass else: + + fn_name = _ALIAS_TO_DISPLAY_FN[line_info.alias].__name__ + node.value = ast.Call( - func=_parse_expr('ecolab.auto_display_utils._display_and_return'), + func=_parse_expr(f'ecolab.auto_display_utils.{fn_name}'), args=[node.value], keywords=[], ) + self.lines_recorder.trailing_stmt_line_nums[line_info.line_num] = ( + line_info + ) except Exception as e: code = '\n'.join(self.lines_recorder.last_lines) print(f'Error for code:\n-----\n{code}\n-----') @@ -182,6 +209,19 @@ def _maybe_display( raise return node + def _is_alias_stmt(self, node: ast.AST) -> bool: + match node: + case ast.Expr(value=ast.Name(id=name)): + pass + case _: + return False + if name not in _ALIAS_TO_DISPLAY_FN: + return False + # The alias is not in the same line as a trailing `;` + if node.end_lineno - 1 not in self.lines_recorder.trailing_stmt_line_nums: + return False + return True + # pylint: disable=invalid-name visit_Assign = _maybe_display visit_AnnAssign = _maybe_display @@ -194,46 +234,100 @@ def _parse_expr(code: str) -> ast.AST: return ast.parse(code, mode='eval').body +@dataclasses.dataclass(frozen=True) +class _LineInfo: + has_trailing: bool + alias: str + line_num: int + + def _has_trailing_semicolon( code_lines: list[str], node: ast.AST, -) -> tuple[bool, bool]: +) -> _LineInfo: """Check if `node` has trailing `;`.""" if isinstance(node, ast.AnnAssign) and node.value is None: - return False, False # `AnnAssign().value` can be `None` (`a: int`) + return _LineInfo( + has_trailing=False, + alias='', + line_num=-1, + ) # `AnnAssign().value` can be `None` (`a: int`) + # Extract the lines of the statement - last_line = code_lines[node.end_lineno - 1] # lineno starts at `1` - # Check if the last character is a `;` token - has_trailing = False + line_num = node.end_lineno - 1 + last_line = code_lines[line_num] # lineno starts at `1` # `node.end_col_offset` is in bytes, so UTF-8 characters count 3. last_part_of_line = last_line.encode('utf-8') last_part_of_line = last_part_of_line[node.end_col_offset :] last_part_of_line = last_part_of_line.decode('utf-8') - for char in last_part_of_line: - if char == ';': - has_trailing = True - elif char == ' ': - continue - elif char == '#': # Comment,... - break - else: # e.g. `a=1;b=2` - has_trailing = False - break - - is_last_statement = True # Assume statement is the last one - for line in code_lines[node.end_lineno :]: # Next statements are all empty - line = line.strip() - if line and not line.startswith('#'): - is_last_statement = False - break - if last_line.startswith(' '): - # statement is inside `if` / `with` / ... - is_last_statement = False - return has_trailing, is_last_statement + + # Check if the last character is a `;` token + has_trailing = False + alias = '' + if match := _detect_trailing_regex().match(last_part_of_line): + has_trailing = True + if match.group(1): + alias = match.group(1) + + return _LineInfo( + has_trailing=has_trailing, + alias=alias, + line_num=line_num, + ) + + +@functools.cache +def _detect_trailing_regex() -> re.Pattern[str]: + """Check if the last character is a `;` token.""" + # Match: + # * `; a` + # * `; a # Some comment` + # * `; # Some comment` + # Do not match: + # * `; a; b` + # * `; a=1` + available_chars = ''.join(_ALIAS_TO_DISPLAY_FN) + return re.compile(f' *; *([{available_chars}])? *(?:#.*)?$') def _display_and_return(x: _T) -> _T: """Print `x` and return `x`.""" IPython.display.display(x) return x + + +def _display_specs_and_return(x: _T) -> _T: + """Print `x` and return `x`.""" + IPython.display.display(etree.spec_like(x)) + return x + + +def _inspect_and_return(x: _T) -> _T: + """Print `x` and return `x`.""" + inspects.inspect(x) + return x + + +def _display_array_and_return(x: _T) -> _T: + """Print `x` and return `x`.""" + html = array_as_img.array_repr_html(x) + if html is None: + IPython.display.display(x) + else: + IPython.display.display(IPython.display.HTML(html)) + return x + + +def _return_quietly(x: _T) -> _T: + """Return `x` without display.""" + return x + + +_ALIAS_TO_DISPLAY_FN = { + '': _display_and_return, + 's': _display_specs_and_return, + 'i': _inspect_and_return, + 'a': _display_array_and_return, + 'q': _return_quietly, +} diff --git a/etils/ecolab/auto_display_utils_test.py b/etils/ecolab/auto_display_utils_test.py index ebe56bad..76666269 100644 --- a/etils/ecolab/auto_display_utils_test.py +++ b/etils/ecolab/auto_display_utils_test.py @@ -15,7 +15,6 @@ """Test.""" import ast -import textwrap from etils import epy from etils.ecolab import auto_display_utils @@ -49,9 +48,9 @@ def test_parsing(): node = ast.parse(code) for stmt in node.body: - assert auto_display_utils._has_trailing_semicolon(code.splitlines(), stmt)[ - 0 - ] + assert auto_display_utils._has_trailing_semicolon( + code.splitlines(), stmt + ).has_trailing code = epy.dedent(""" a=1 @@ -64,6 +63,8 @@ def test_parsing(): 1, 3,) # asdasd a=1;b=2 + a=1;iii + a=1 # ;a d = 3 # AAA ( 1, @@ -83,67 +84,43 @@ def test_parsing(): for stmt in node.body: assert not auto_display_utils._has_trailing_semicolon( code.splitlines(), stmt - )[0], f"Error for: `{ast.unparse(stmt)}`" - - code = epy.dedent(""" - a=1; - b=2; - """) - node = ast.parse(code) - - assert auto_display_utils._has_trailing_semicolon( - code.splitlines(), node.body[0] - ) == (True, False) - assert auto_display_utils._has_trailing_semicolon( - code.splitlines(), node.body[1] - ) == (True, True) + ).has_trailing, f'Error for: `{ast.unparse(stmt)}`' @pytest.mark.parametrize( - "code", + 'code, alias', [ - """ - b = 2 - other = 123; - - a = 1; - """, - """ - a=1; - - # Some comment - # other comment - """, - """ - a=1; - - - """, - "a=1; # Comment", + ('x;', ''), + ('x;#i', ''), + ('x;a', 'a'), + ('x;i', 'i'), + ('x ;i', 'i'), + ('x ; i', 'i'), + ('x; i', 'i'), + ('x; i ', 'i'), + ('x;i#Comment', 'i'), + ('x;i #Comment', 'i'), + ('x; i # Comment', 'i'), ], ) -def test_ignore_last(code): - code = textwrap.dedent(code) +def test_alias(code: str, alias: str): node = ast.parse(code) - - for stmt in node.body[:-1]: - assert not auto_display_utils._has_trailing_semicolon( - code.splitlines(), stmt - )[1] - assert auto_display_utils._has_trailing_semicolon( - code.splitlines(), node.body[-1] - )[1] + info = auto_display_utils._has_trailing_semicolon( + code.splitlines(), node.body[0] + ) + assert info.has_trailing + assert info.alias == alias -def test_ignore_last_not_last(): - code = textwrap.dedent(""" - if x: - a = 1; - """) +@pytest.mark.parametrize( + 'code', + [ + 'x;aa', + 'x;a=1', + ], +) +def test_noalias(code): node = ast.parse(code) - - has_trailing, is_last_statement = auto_display_utils._has_trailing_semicolon( - code.splitlines(), node.body[-1].body[-1] # pytype: disable=attribute-error - ) - assert has_trailing - assert not is_last_statement + assert not auto_display_utils._has_trailing_semicolon( + code.splitlines(), node.body[0] + ).has_trailing diff --git a/pyproject.toml b/pyproject.toml index f4aacce6..024f15c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ name = "etils" description = "Collection of common python utils" readme = "README.md" -requires-python = ">=3.9" +requires-python = ">=3.10" license = {file = "LICENSE"} authors = [{name = "Conchylicultor", email="etils@google.com"}] classifiers = [ @@ -48,6 +48,7 @@ ecolab = [ "protobuf", "etils[enp]", "etils[epy]", + "etils[etree]", ] edc = [ # Do not add anything here. `edc` is an alias for `epy`