!51 bug(tables): [#102] Probely CLI - get commands table's header shows up funny Closes #102

Author: João Silva

probely/cli/commands/findings/get.py

@@ -2,8 +2,9 @@
 
 from probely.cli.commands.findings.schemas import FindingsApiFiltersSchema
 from probely.cli.common import prepare_filters_for_api
-from probely.cli.enums import EntityTypeEnum, OutputEnum
+from probely.cli.enums import OutputEnum
 from probely.cli.renderers import OutputRenderer
+from probely.cli.tables.finding_table import FindingTable
 from probely.exceptions import ProbelyCLIValidation
 from probely.sdk.findings import list_findings, retrieve_findings
 
@@ -23,6 +24,6 @@ def findings_get_command_handler(args: argparse.Namespace):
         records=findings_generator,
         output_type=output_type,
         console=args.console,
-        entity_type=EntityTypeEnum.FINDING,
+        table_cls=FindingTable,
     )
     renderer.render()

probely/cli/commands/scans/get.py

@@ -2,8 +2,9 @@
 
 from probely.cli.commands.scans.schemas import ScanApiFiltersSchema
 from probely.cli.common import prepare_filters_for_api
-from probely.cli.enums import EntityTypeEnum, OutputEnum
+from probely.cli.enums import OutputEnum
 from probely.cli.renderers import OutputRenderer
+from probely.cli.tables.scan_table import ScanTable
 from probely.exceptions import ProbelyCLIValidation
 from probely.sdk.scans import list_scans, retrieve_scans
 
@@ -25,6 +26,6 @@ def scans_get_command_handler(args: argparse.Namespace):
         records=scans_generator,
         output_type=output_type,
         console=args.console,
-        entity_type=EntityTypeEnum.SCAN,
+        table_cls=ScanTable,
     )
     renderer.render()

probely/cli/commands/targets/get.py

@@ -1,8 +1,9 @@
 from probely.cli.commands.targets.schemas import TargetApiFiltersSchema
 from probely.cli.common import prepare_filters_for_api
-from probely.cli.enums import EntityTypeEnum, OutputEnum
+from probely.cli.enums import OutputEnum
 
 from probely.cli.renderers import OutputRenderer
+from probely.cli.tables.targets_table import TargetTable
 from probely.exceptions import ProbelyCLIValidation
 from probely.sdk.targets import list_targets, retrieve_targets
 
@@ -28,6 +29,6 @@ def targets_get_command_handler(args):
         records=targets_generator,
         output_type=output_type,
         console=args.console,
-        entity_type=EntityTypeEnum.TARGET,
+        table_cls=TargetTable,
     )
     renderer.render()

probely/cli/enums.py

@@ -1,14 +1,6 @@
-from enum import Enum
-
 from probely.sdk.enums import ProbelyCLIEnum
 
 
-class EntityTypeEnum(Enum):
-    SCAN = "scan"
-    FINDING = "finding"
-    TARGET = "target"
-
-
 class OutputEnum(ProbelyCLIEnum):
     YAML = "yaml"
     JSON = "json"

probely/cli/renderers.py

@@ -13,9 +13,9 @@
 import yaml
 from dateutil import parser
 from rich.console import Console
-from rich.live import Live
 
-from probely.cli.enums import EntityTypeEnum, OutputEnum
+from probely.cli.enums import OutputEnum
+from probely.cli.tables.base_table import BaseOutputTable
 from probely.sdk.enums import ProbelyCLIEnum
 
 UNKNOWN_VALUE_REP = "UNKNOWN"
@@ -32,12 +32,12 @@ def __init__(
         records: Generator[dict, None, None],
         output_type: Optional[OutputEnum],
         console: Console,
-        entity_type: EntityTypeEnum,
+        table_cls: Type[BaseOutputTable],
     ):
         self.records = records
         self.output_type = output_type
         self.console = console
-        self.entity_type = entity_type
+        self.table_cls = table_cls
 
     def render(self) -> None:
         if self.output_type == OutputEnum.JSON:
@@ -62,15 +62,13 @@ def _render_yaml(self) -> None:
             self.console.print(yaml.dump([record], indent=2, width=sys.maxsize))
 
     def _render_table(self) -> None:
-        from probely.cli.tables.table_factory import (  # Avoid circular import
-            TableFactory,
-        )
-
-        table_cls = TableFactory.get_table_class(self.entity_type)
-        table = table_cls().create_table(show_header=True)
-        with Live(table, console=self.console):
-            for record in self.records:
-                table_cls().add_row(table, record)
+        table = self.table_cls.create_table(show_header=True)
+        self.console.print(table)
+
+        for record in self.records:
+            table = self.table_cls.create_table(show_header=False)
+            self.table_cls.add_row(table, record)
+            self.console.print(table)
 
 
 def get_printable_enum_value(enum: Type[ProbelyCLIEnum], api_enum_value: str) -> str:

probely/cli/tables/base_table.py

@@ -4,15 +4,17 @@
 
 
 class BaseOutputTable(ABC):
+    @classmethod
     @abstractmethod
-    def create_table(self, show_header: bool) -> Table:
+    def create_table(cls, show_header: bool) -> Table:
         """
         Initializes and returns a Rich Table with predefined columns.
         """
         pass
 
+    @classmethod
     @abstractmethod
-    def add_row(self, table: Table, record: dict) -> None:
+    def add_row(cls, table: Table, record: dict) -> None:
         """
         Adds a single row to the provided Rich Table based on the record data.
         """

probely/cli/tables/finding_table.py

@@ -1,4 +1,5 @@
 from typing import Dict
+
 from rich.table import Table
 
 from probely.cli.renderers import (
@@ -13,7 +14,8 @@
 
 
 class FindingTable(BaseOutputTable):
-    def create_table(self, show_header: bool = False) -> Table:
+    @classmethod
+    def create_table(cls, show_header: bool = False) -> Table:
         table = Table(show_header=show_header, box=None)
 
         table.add_column("ID", width=18)
@@ -26,7 +28,8 @@ def create_table(self, show_header: bool = False) -> Table:
 
         return table
 
-    def add_row(self, table: Table, finding: Dict) -> None:
+    @classmethod
+    def add_row(cls, table: Table, finding: Dict) -> None:
         target = finding.get("target", {})
 
         table.add_row(

probely/cli/tables/scan_table.py

@@ -5,7 +5,8 @@
 
 
 class ScanTable(BaseOutputTable):
-    def create_table(self, show_header: bool = False) -> Table:
+    @classmethod
+    def create_table(cls, show_header: bool = False) -> Table:
         table = Table(show_header=show_header, box=None)
 
         table.add_column("ID", width=12)
@@ -20,7 +21,8 @@ def create_table(self, show_header: bool = False) -> Table:
 
         return table
 
-    def add_row(self, table: Table, scan: dict) -> None:
+    @classmethod
+    def add_row(cls, table: Table, scan: dict) -> None:
         target = scan.get("target", {})
         site = target.get("site", {})
 

probely/cli/tables/table_factory.py

@@ -1,6 +1,7 @@
 from typing import Dict
 
 from rich.table import Table
+
 from probely.cli.renderers import (
     get_printable_enum_value,
     get_printable_labels,
@@ -11,7 +12,8 @@
 
 
 class TargetTable(BaseOutputTable):
-    def create_table(self, show_header: bool = False) -> Table:
+    @classmethod
+    def create_table(cls, show_header: bool = False) -> Table:
         table = Table(show_header=show_header, box=None)
 
         table.add_column("ID", width=12)
@@ -23,7 +25,8 @@ def create_table(self, show_header: bool = False) -> Table:
 
         return table
 
-    def add_row(self, table: Table, target: Dict) -> None:
+    @classmethod
+    def add_row(cls, table: Table, target: Dict) -> None:
         site = target.get("site")
 
         table.add_row(