From f4402ae76bdab12649d80726195449bcb8d92f7e Mon Sep 17 00:00:00 2001
From: LeaveMyYard <zhukovpavel2001@gmail.com>
Date: Mon, 23 Oct 2023 11:25:39 +0300
Subject: [PATCH] Add --width parameter for the console

---
 robusta_krr/core/models/config.py | 19 ++++++++++++++++---
 robusta_krr/main.py               |  4 ++++
 robusta_krr/utils/print.py        | 11 ++---------
 3 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/robusta_krr/core/models/config.py b/robusta_krr/core/models/config.py
index 3d6aabc0..5a26c490 100644
--- a/robusta_krr/core/models/config.py
+++ b/robusta_krr/core/models/config.py
@@ -53,6 +53,7 @@ class Config(pd.BaseSettings):
     format: str
     strategy: str
     log_to_stderr: bool
+    width: Optional[int] = pd.Field(None, ge=1)
 
     # Outputs Settings
     file_output: Optional[str] = pd.Field(None)
@@ -62,11 +63,11 @@ class Config(pd.BaseSettings):
 
     # Internal
     inside_cluster: bool = False
-    console: Optional[Console] = None
+    _logging_console: Optional[Console] = pd.PrivateAttr(None)
+    _result_console: Optional[Console] = pd.PrivateAttr(None)
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
-        self.console = Console(stderr=self.log_to_stderr)
 
     @property
     def Formatter(self) -> formatters.FormatterFunc:
@@ -112,6 +113,18 @@ def validate_format(cls, v: str) -> str:
     def context(self) -> Optional[str]:
         return self.clusters[0] if self.clusters != "*" and self.clusters else None
 
+    @property
+    def logging_console(self) -> Console:
+        if getattr(self, "_logging_console") is None:
+            self._logging_console = Console(file=sys.stderr if self.log_to_stderr else sys.stdout, width=self.width)
+        return self._logging_console
+
+    @property
+    def result_console(self) -> Console:
+        if getattr(self, "_result_console") is None:
+            self._result_console = Console(file=sys.stdout, width=self.width)
+        return self._result_console
+
     def load_kubeconfig(self) -> None:
         try:
             config.load_incluster_config()
@@ -130,7 +143,7 @@ def set_config(config: Config) -> None:
             level="NOTSET",
             format="%(message)s",
             datefmt="[%X]",
-            handlers=[RichHandler(console=Console(file=sys.stderr if settings.log_to_stderr else sys.stdout))],
+            handlers=[RichHandler(console=config.logging_console)],
         )
         logging.getLogger("").setLevel(logging.CRITICAL)
         logger.setLevel(logging.DEBUG if config.verbose else logging.CRITICAL if config.quiet else logging.INFO)
diff --git a/robusta_krr/main.py b/robusta_krr/main.py
index 9710c47a..8c19f1bf 100644
--- a/robusta_krr/main.py
+++ b/robusta_krr/main.py
@@ -205,6 +205,9 @@ def run_strategy(
                 log_to_stderr: bool = typer.Option(
                     False, "--logtostderr", help="Pass logs to stderr", rich_help_panel="Logging Settings"
                 ),
+                width: Optional[int] = typer.Option(
+                    None, "--width", help="Width of the output. Will use console width by default.", rich_help_panel="Logging Settings"
+                ),
                 file_output: Optional[str] = typer.Option(
                     None, "--fileoutput", help="Print the output to a file", rich_help_panel="Output Settings"
                 ),
@@ -245,6 +248,7 @@ def run_strategy(
                         memory_min_value=memory_min_value,
                         quiet=quiet,
                         log_to_stderr=log_to_stderr,
+                        width=width,
                         file_output=file_output,
                         slack_output=slack_output,
                         strategy=_strategy_name,
diff --git a/robusta_krr/utils/print.py b/robusta_krr/utils/print.py
index 1963c3e0..093cc445 100644
--- a/robusta_krr/utils/print.py
+++ b/robusta_krr/utils/print.py
@@ -1,21 +1,14 @@
-import sys
-
-from rich import print as r_print
-
 from robusta_krr.core.models.config import settings
 
-py_print = print
-
 
 def print(*objects, rich: bool = True, force: bool = False) -> None:
     """
     A wrapper around `rich.print` that prints only if `settings.quiet` is False.
     """
-    print_func = r_print if rich else py_print
-    output = sys.stdout if force or not settings.log_to_stderr else sys.stderr
+    print_func = settings.logging_console.print if rich else print
 
     if not settings.quiet or force:
-        print_func(*objects, file=output)  # type: ignore
+        print_func(*objects)  # type: ignore
 
 
 __all__ = ["print"]