From 5cd64f1736b4bd3843b2edec289c8e75af650e8b Mon Sep 17 00:00:00 2001
From: JM Lopez <jm.lopez@utoronto.ca>
Date: Wed, 16 Oct 2019 12:19:16 -0400
Subject: [PATCH] Allows setting maximum column width --columns-width - If
 specified --columns-width or HASS_COL_WIDTH env var   truncates column values
 - Uses default const.COLUMNS_WIDTH_STR - Fixes #253

---
 homeassistant_cli/cli.py                       | 13 +++++++++++++
 homeassistant_cli/config.py                    |  3 ++-
 homeassistant_cli/const.py                     |  1 +
 homeassistant_cli/helper.py                    | 18 +++++++++++++++++-
 .../basic_entities_table_fixed_width.txt       |  4 ++++
 tests/test_state.py                            | 11 +++++++++++
 6 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 tests/fixtures/basic_entities_table_fixed_width.txt

diff --git a/homeassistant_cli/cli.py b/homeassistant_cli/cli.py
index 3d9f826..ff9012c 100644
--- a/homeassistant_cli/cli.py
+++ b/homeassistant_cli/cli.py
@@ -188,6 +188,17 @@ def _default_token() -> Optional[str]:
         ' Example: ENTITY=entity_id, NAME=attributes.friendly_name'
     ),
 )
+@click.option(
+    '--columns-width',
+    default=None,
+    type=click.INT,
+    envvar='HASS_COL_WIDTH',
+    show_envvar=True,
+    help=(
+        'Columns custom width. '
+        'If specified truncates column values (default: auto)'
+    ),
+)
 @click.option(
     '--no-headers',
     default=False,
@@ -220,6 +231,7 @@ def cli(
     showexceptions: bool,
     cert: str,
     columns: str,
+    columns_width: int,
     no_headers: bool,
     table_format: str,
     sort_by: Optional[str],
@@ -236,6 +248,7 @@ def cli(
     ctx.showexceptions = showexceptions
     ctx.cert = cert
     ctx.columns = to_tuples(columns)
+    ctx.columns_width = columns_width
     ctx.no_headers = no_headers
     ctx.table_format = table_format
     ctx.sort_by = sort_by  # type: ignore
diff --git a/homeassistant_cli/config.py b/homeassistant_cli/config.py
index 747a7be..43d5cb6 100644
--- a/homeassistant_cli/config.py
+++ b/homeassistant_cli/config.py
@@ -38,7 +38,7 @@ def _locate_ha() -> Optional[str]:
     listener = _ZeroconfListener()
     zeroconf.ServiceBrowser(_zeroconf, "_home-assistant._tcp.local.", listener)
     try:
-        import time
+        import time  # pylint: disable=import-outside-toplevel
 
         retries = 0
         while not listener.services and retries < 5:
@@ -120,6 +120,7 @@ def __init__(self) -> None:
         self.session = None  # type: Optional[Session]
         self.cert = None  # type: Optional[str]
         self.columns = None  # type: Optional[List[Tuple[str, str]]]
+        self.columns_width = None  # type: Optional[int]
         self.no_headers = False
         self.table_format = 'plain'
         self.sort_by = None
diff --git a/homeassistant_cli/const.py b/homeassistant_cli/const.py
index 0147ecc..1827e58 100644
--- a/homeassistant_cli/const.py
+++ b/homeassistant_cli/const.py
@@ -19,3 +19,4 @@
     ('CHANGED', 'last_changed'),
 ]
 COLUMNS_SERVICES = [('DOMAIN', 'domain'), ("SERVICE", "domain.services[*]")]
+COLUMNS_WIDTH_STR = "..."
diff --git a/homeassistant_cli/helper.py b/homeassistant_cli/helper.py
index 16e66ee..14cdaa2 100644
--- a/homeassistant_cli/helper.py
+++ b/homeassistant_cli/helper.py
@@ -51,6 +51,7 @@ def raw_format_output(
     data: Union[Dict[str, Any], List[Dict[str, Any]]],
     yamlparser: YAML,
     columns: Optional[List] = None,
+    columns_width: Optional[int] = None,
     no_headers: bool = False,
     table_format: str = 'plain',
     sort_by: Optional[str] = None,
@@ -99,7 +100,20 @@ def raw_format_output(
                 val = [match.value for match in fmtpair[1].find(item)]
                 row.append(", ".join(map(str, val)))
             result.append(row)
-
+        # Truncates data
+        if columns_width:
+            max_str = columns_width - len(const.COLUMNS_WIDTH_STR)
+            result = [
+                [
+                    (
+                        c
+                        if len(c) < columns_width
+                        else c[:max_str] + const.COLUMNS_WIDTH_STR
+                    )
+                    for c in row
+                ]
+                for row in result
+            ]
         res = tabulate(
             result, headers=headers, tablefmt=table_format
         )  # type: str
@@ -130,6 +144,7 @@ def format_output(
     ctx: Configuration,
     data: List[Dict[str, Any]],
     columns: Optional[List] = None,
+    columns_width: Optional[int] = None,
 ) -> str:
     """Format data to output based on settings in ctx/Context."""
     return raw_format_output(
@@ -137,6 +152,7 @@ def format_output(
         data,
         ctx.yaml(),
         columns,
+        ctx.columns_width,
         ctx.no_headers,
         ctx.table_format,
         ctx.sort_by,
diff --git a/tests/fixtures/basic_entities_table_fixed_width.txt b/tests/fixtures/basic_entities_table_fixed_width.txt
new file mode 100644
index 0000000..d9e0a7b
--- /dev/null
+++ b/tests/fixtures/basic_entities_table_fixed_width.txt
@@ -0,0 +1,4 @@
+ENTITY        DESCRIPTION      STATE    CHANGED
+sensor.one    friendly lon...  on       2018-12-02T1...
+sensor.two                     off      2018-12-01T1...
+sensor.three                   off      2018-12-03T1...
diff --git a/tests/test_state.py b/tests/test_state.py
index 1ba9e35..4bcb601 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -137,6 +137,17 @@ def test_state_list_table_columns_sortby(
     )
 
 
+def test_state_list_table_columns_width(
+    basic_entities_text, basic_entities_table_fixed_width_text
+) -> None:
+    """Test table columns."""
+    output_formats(
+        ["--output=table", "--columns-width=15", "state", "list"],
+        basic_entities_text,
+        basic_entities_table_fixed_width_text,
+    )
+
+
 def test_state_list_no_header(
     basic_entities_text, basic_entities_table_no_header_text
 ) -> None: