diff --git a/kioss/_execution/_core.py b/kioss/_execution/_core.py index 1df2372..6c05735 100644 --- a/kioss/_execution/_core.py +++ b/kioss/_execution/_core.py @@ -52,9 +52,10 @@ def __next__(self) -> R: class LoggingIteratorWrapper(IteratorWrapper[T]): - def __init__(self, iterator: Iterator[T], what: str) -> None: + def __init__(self, iterator: Iterator[T], what: str, colored: bool) -> None: super().__init__(iterator) self.what = what + self.colored = colored self.yields_count = 0 self.errors_count = 0 self.last_log_at_yields_count: Optional[int] = None @@ -63,11 +64,16 @@ def __init__(self, iterator: Iterator[T], what: str) -> None: def _log(self) -> None: errors_summary = f"with {self.errors_count} error{'s' if self.errors_count > 1 else ''} produced" - if self.errors_count > 0: + if self.colored and self.errors_count > 0: errors_summary = _util.bold(_util.colorize_in_red(errors_summary)) - yields_summary = _util.bold( - f"{self.yields_count} {self.what[:-1] if self.yields_count <= 1 and self.what.endswith('s') else self.what} have been yielded" + what = ( + self.what[:-1] + if self.yields_count <= 1 and self.what.endswith("s") + else self.what ) + yields_summary = f"{self.yields_count} {what} have been yielded" + if self.colored: + yields_summary = _util.bold(yields_summary) elapsed_time = f"in elapsed time {datetime.fromtimestamp(time.time()) - datetime.fromtimestamp(self.start_time)}" _util.LOGGER.info("%s, %s, %s", yields_summary, elapsed_time, errors_summary) diff --git a/kioss/_pipe.py b/kioss/_pipe.py index 2f49036..87f2244 100644 --- a/kioss/_pipe.py +++ b/kioss/_pipe.py @@ -38,13 +38,10 @@ def __iter__(self) -> Iterator[T]: def __add__(self, other: "APipe[T]") -> "APipe[T]": return self.chain(other) - def __repr__(self) -> str: - """ - Explain the plan of the pipe - """ + def explain(self, colored: bool = False) -> str: if EXPLAINING_VISITOR_CLASS is None: raise ValueError("_pipe.EXPLAINING_VISITOR_CLASS is None") - return self._accept(EXPLAINING_VISITOR_CLASS()) + return self._accept(EXPLAINING_VISITOR_CLASS(colored)) @abstractmethod def _accept(self, visitor: "AVisitor") -> Any: @@ -175,17 +172,18 @@ def catch( """ return CatchPipe(self, *classes, when=when) - def log(self, what: str = "elements") -> "APipe[T]": + def log(self, what: str = "elements", colored: bool = False) -> "APipe[T]": """ Log the elements of the Pipe as they are iterated. Args: what (str): name the objects yielded by the pipe for clearer logs, must be a plural descriptor. + colored (bool): whether or not to use ascii colorization. Returns: Pipe[T]: A new Pipe instance with logging capability. """ - return LogPipe(self, what) + return LogPipe(self, what, colored) def collect(self, n_samples: Optional[int] = None) -> List[T]: """ @@ -239,7 +237,7 @@ def register_error_sample(error): return True safe_plan = plan.catch(Exception, when=register_error_sample) - _util.LOGGER.info(repr(safe_plan)) + _util.LOGGER.info(safe_plan.explain(colored=False)) samples = safe_plan.collect(n_samples=n_samples) if errors_count > 0: _util.LOGGER.error( @@ -305,9 +303,10 @@ def _accept(self, visitor: "AVisitor") -> Any: class LogPipe(APipe[T]): - def __init__(self, upstream: APipe[T], what: str = "elements"): + def __init__(self, upstream: APipe[T], what: str, colored: bool): self.upstream: APipe[T] = upstream self.what = what + self.colored = colored def _accept(self, visitor: "AVisitor") -> Any: return visitor.visit_log_pipe(self) diff --git a/kioss/_visit/_explanation.py b/kioss/_visit/_explanation.py index bf663a3..221c21f 100644 --- a/kioss/_visit/_explanation.py +++ b/kioss/_visit/_explanation.py @@ -7,21 +7,34 @@ class ExplainingVisitor(AVisitor): HEADER = "Pipe's plan:" - def __init__(self, initial_margin: int = 0, add_header: bool = True): + def __init__(self, colored: bool, initial_margin: int = 0, add_header: bool = True): + self.colored = colored self.current_margin = initial_margin self.margin_step = 2 self.add_header = add_header def additional_explain_lines(self, name: str, descr: str) -> str: - return f"{' '*self.current_margin}{_util.colorize_in_grey('└' + '─'*(self.margin_step - 1))}•{_util.colorize_in_red(name)}({descr})\n" + margin = " " * self.current_margin + if self.add_header: + linking_symbols = " " * self.margin_step + "•" + else: + linking_symbols = "└" + "─" * (self.margin_step - 1) + "•" + + if self.colored: + linking_symbols = _util.colorize_in_grey(linking_symbols) + name = _util.colorize_in_red(name) + return f"{margin}{linking_symbols}{name}({descr})\n" def visit_any_pipe(self, pipe: _pipe.APipe, name: str, descr: str) -> str: + additional_explain_lines = self.additional_explain_lines(name, descr) if self.add_header: - header = _util.bold(ExplainingVisitor.HEADER) + "\n" + if self.colored: + header = _util.bold(ExplainingVisitor.HEADER) + "\n" + else: + header = ExplainingVisitor.HEADER + "\n" self.add_header = False else: header = "" - additional_explain_lines = self.additional_explain_lines(name, descr) self.current_margin += self.margin_step if pipe.upstream is not None: upstream_repr = pipe.upstream._accept(self) @@ -37,7 +50,9 @@ def visit_chain_pipe(self, pipe: _pipe.ChainPipe) -> Any: chained_pipes_repr = "".join( map( lambda pipe: pipe._accept( - ExplainingVisitor(self.current_margin, add_header=False) + ExplainingVisitor( + self.colored, self.current_margin, add_header=False + ) ), pipe.others, ) diff --git a/kioss/_visit/_iter_production.py b/kioss/_visit/_iter_production.py index 413462d..5e2cfc5 100644 --- a/kioss/_visit/_iter_production.py +++ b/kioss/_visit/_iter_production.py @@ -68,4 +68,6 @@ def visit_catch_pipe(self, pipe: _pipe.CatchPipe[T]) -> Iterator[T]: ) def visit_log_pipe(self, pipe: _pipe.LogPipe[T]) -> Iterator[T]: - return _core.LoggingIteratorWrapper(pipe.upstream._accept(self), pipe.what) + return _core.LoggingIteratorWrapper( + pipe.upstream._accept(self), pipe.what, pipe.colored + ) diff --git a/setup.py b/setup.py index 785ddc5..80e76db 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ setup( name='kioss', - version='0.6.4', + version='0.6.5', packages=find_packages(), package_data={"kioss": ["py.typed"]}, url='http://github.com/bonnal-enzo/kioss', diff --git a/tests/test_legacy.py b/tests/test_legacy.py index 35d3180..b878b88 100644 --- a/tests/test_legacy.py +++ b/tests/test_legacy.py @@ -465,6 +465,15 @@ def test_log(self) -> None: [[1, 2], [3]], ) + ( + Pipe("12-3".__iter__) + .log("chars") + .map(int) + .log("ints", colored=True) + .batch(2) + .superintend(raise_if_more_errors_than=1) + ) + def test_partial_iteration(self) -> None: first_elem = next( iter( @@ -522,7 +531,7 @@ def test_generator_already_generating(self) -> None: Counter(list(range(N)) + list(range(N)) + list(range(N))), ) - def test_repr(self) -> None: + def test_explain(self) -> None: p = ( Pipe(range(8).__iter__) .filter(lambda _: True) @@ -538,9 +547,10 @@ def test_repr(self) -> None: ) .catch(ValueError, TypeError, when=lambda e: True) ) - a = repr(p) + a = p.explain() p.collect() - b = repr(p) + b = p.explain() + c = p.explain(colored=True) self.assertEqual(a, b) - self.assertGreater(len(a), 32) - print(a) + self.assertGreater(len(c), len(a)) + print(c)