From c02ffac403211f5cda7c3023d9d9333f9e805a71 Mon Sep 17 00:00:00 2001
From: onyx-and-iris <code@onyxandiris.online>
Date: Sun, 22 Jun 2025 02:52:27 +0100
Subject: [PATCH] replace terminaltables with rich tables.

allow rich to handle all console output.

util.check_mark is now used to pass back colourless check/cross marks if NO_COLOR is set or --style/SLOBS_STYLE was not set.
---
 src/slobs_cli/audio.py           | 64 +++++++++++++++---------
 src/slobs_cli/console.py         | 21 ++++++++
 src/slobs_cli/errors.py          | 12 ++---
 src/slobs_cli/record.py          | 13 ++---
 src/slobs_cli/replaybuffer.py    | 11 ++--
 src/slobs_cli/scene.py           | 86 +++++++++++++++++++-------------
 src/slobs_cli/scenecollection.py | 77 +++++++++++++++++-----------
 src/slobs_cli/stream.py          | 13 ++---
 src/slobs_cli/studiomode.py      | 15 +++---
 src/slobs_cli/util.py            | 15 ++++++
 10 files changed, 209 insertions(+), 118 deletions(-)
 create mode 100644 src/slobs_cli/console.py
 create mode 100644 src/slobs_cli/util.py

diff --git a/src/slobs_cli/audio.py b/src/slobs_cli/audio.py
index 9758975..b743d15 100644
--- a/src/slobs_cli/audio.py
+++ b/src/slobs_cli/audio.py
@@ -3,8 +3,10 @@
 import asyncclick as click
 from anyio import create_task_group
 from pyslobs import AudioService
-from terminaltables3 import AsciiTable
+from rich.table import Table
+from rich.text import Text
 
+from . import console, util
 from .cli import cli
 from .errors import SlobsCliError
 
@@ -25,32 +27,40 @@ async def list(ctx: click.Context, id: bool = False):
     async def _run():
         sources = await as_.get_sources()
         if not sources:
-            click.echo('No audio sources found.')
+            console.out.print('No audio sources found.')
             conn.close()
             return
 
-        table_data = [
-            ['Audio Device Name', 'ID', 'Muted']
-            if id
-            else ['Audio Device Name', 'Muted']
-        ]
+        style = ctx.obj['style']
+        table = Table(
+            show_header=True, header_style=style.header, border_style=style.border
+        )
+
+        if id:
+            columns = [
+                ('Audio Source Name', 'left'),
+                ('Muted', 'center'),
+                ('ID', 'left'),
+            ]
+        else:
+            columns = [
+                ('Audio Source Name', 'left'),
+                ('Muted', 'center'),
+            ]
+        for col_name, col_justify in columns:
+            table.add_column(Text(col_name, justify='center'), justify=col_justify)
+
         for source in sources:
             model = await source.get_model()
 
-            to_append = [click.style(model.name, fg='blue')]
+            to_append = [Text(model.name, style=style.cell)]
+            to_append.append(util.check_mark(ctx, model.muted))
             if id:
-                to_append.append(model.source_id)
-            to_append.append('✅' if model.muted else '❌')
+                to_append.append(Text(model.source_id, style=style.cell))
 
-            table_data.append(to_append)
+            table.add_row(*to_append)
 
-        table = AsciiTable(table_data)
-        table.justify_columns = {
-            0: 'left',
-            1: 'left' if id else 'center',
-            2: 'center' if id else None,
-        }
-        click.echo(table.table)
+        console.out.print(table)
 
         conn.close()
 
@@ -78,7 +88,7 @@ async def mute(ctx: click.Context, source_name: str):
             raise SlobsCliError(f'Audio source "{source_name}" not found.')
 
         await source.set_muted(True)
-        click.echo(f'{source_name} muted successfully.')
+        console.out.print(f'{console.highlight(ctx, source_name)} muted successfully.')
         conn.close()
 
     try:
@@ -109,7 +119,9 @@ async def unmute(ctx: click.Context, source_name: str):
             raise SlobsCliError(f'Audio source "{source_name}" not found.')
 
         await source.set_muted(False)
-        click.echo(f'{source_name} unmuted successfully.')
+        console.out.print(
+            f'{console.highlight(ctx, source_name)} unmuted successfully.'
+        )
         conn.close()
 
     try:
@@ -136,10 +148,14 @@ async def toggle(ctx: click.Context, source_name: str):
             if model.name.lower() == source_name.lower():
                 if model.muted:
                     await source.set_muted(False)
-                    click.echo(f'{source_name} unmuted successfully.')
+                    console.out.print(
+                        f'{console.highlight(ctx, source_name)} unmuted successfully.'
+                    )
                 else:
                     await source.set_muted(True)
-                    click.echo(f'{source_name} muted successfully.')
+                    console.out.print(
+                        f'{console.highlight(ctx, source_name)} muted successfully.'
+                    )
                 conn.close()
                 break
         else:  # If no source by the given name was found
@@ -168,8 +184,8 @@ async def status(ctx: click.Context, source_name: str):
         for source in sources:
             model = await source.get_model()
             if model.name.lower() == source_name.lower():
-                click.echo(
-                    f'"{source_name}" is {"muted" if model.muted else "unmuted"}.'
+                console.out.print(
+                    f'{console.highlight(ctx, source_name)} is {"muted" if model.muted else "unmuted"}.'
                 )
                 conn.close()
                 return
diff --git a/src/slobs_cli/console.py b/src/slobs_cli/console.py
new file mode 100644
index 0000000..33b2b49
--- /dev/null
+++ b/src/slobs_cli/console.py
@@ -0,0 +1,21 @@
+"""module for console output handling."""
+
+import asyncclick as click
+from rich.console import Console
+
+out = Console()
+err = Console(stderr=True, style='bold red')
+
+
+def highlight(ctx: click.Context, text: str) -> str:
+    """Highlight text for console output."""
+    if ctx.obj['style'].name == 'no_colour':
+        return text
+    return f'[{ctx.obj["style"].highlight}]{text}[/{ctx.obj["style"].highlight}]'
+
+
+def warning(ctx: click.Context, text: str) -> str:
+    """Format warning text for console output."""
+    if ctx.obj['style'].name == 'no_colour':
+        return text
+    return f'[magenta]{text}[/magenta]'
diff --git a/src/slobs_cli/errors.py b/src/slobs_cli/errors.py
index 4d8790f..299eb64 100644
--- a/src/slobs_cli/errors.py
+++ b/src/slobs_cli/errors.py
@@ -4,6 +4,8 @@ import json
 
 import asyncclick as click
 
+from . import console
+
 
 class SlobsCliError(click.ClickException):
     """Base class for all Slobs CLI errors."""
@@ -14,8 +16,8 @@ class SlobsCliError(click.ClickException):
         self.exit_code = 1
 
     def show(self):
-        """Display the error message in red."""
-        click.secho(f'Error: {self.message}', fg='red', err=True)
+        """Display the error message in red and write to stderr."""
+        console.err.print(f'Error: {self.message}')
 
 
 class SlobsCliProtocolError(SlobsCliError):
@@ -36,10 +38,8 @@ class SlobsCliProtocolError(SlobsCliError):
         """Display the protocol error message in red."""
         match self.protocol_code:
             case -32600:
-                click.secho(
-                    'Oops! Looks like we hit a rate limit for this command. Please try again later.',
-                    fg='red',
-                    err=True,
+                console.err.print(
+                    'Oops! Looks like we hit a rate limit for this command. Please try again later.'
                 )
             case _:
                 # Fall back to the base error display for unknown protocol codes
diff --git a/src/slobs_cli/record.py b/src/slobs_cli/record.py
index 5597df5..5442b7e 100644
--- a/src/slobs_cli/record.py
+++ b/src/slobs_cli/record.py
@@ -4,6 +4,7 @@ import asyncclick as click
 from anyio import create_task_group
 from pyslobs import StreamingService
 
+from . import console
 from .cli import cli
 from .errors import SlobsCliError
 
@@ -29,7 +30,7 @@ async def start(ctx: click.Context):
             raise SlobsCliError('Recording is already active.')
 
         await ss.toggle_recording()
-        click.echo('Recording started.')
+        console.out.print('Recording started.')
 
         conn.close()
 
@@ -58,7 +59,7 @@ async def stop(ctx: click.Context):
             raise SlobsCliError('Recording is already inactive.')
 
         await ss.toggle_recording()
-        click.echo('Recording stopped.')
+        console.out.print('Recording stopped.')
 
         conn.close()
 
@@ -83,9 +84,9 @@ async def status(ctx: click.Context):
         active = model.recording_status != 'offline'
 
         if active:
-            click.echo('Recording is currently active.')
+            console.out.print('Recording is currently active.')
         else:
-            click.echo('Recording is currently inactive.')
+            console.out.print('Recording is currently inactive.')
 
         conn.close()
 
@@ -107,9 +108,9 @@ async def toggle(ctx: click.Context):
 
         await ss.toggle_recording()
         if active:
-            click.echo('Recording stopped.')
+            console.out.print('Recording stopped.')
         else:
-            click.echo('Recording started.')
+            console.out.print('Recording started.')
 
         conn.close()
 
diff --git a/src/slobs_cli/replaybuffer.py b/src/slobs_cli/replaybuffer.py
index 65fce19..cf673ec 100644
--- a/src/slobs_cli/replaybuffer.py
+++ b/src/slobs_cli/replaybuffer.py
@@ -4,6 +4,7 @@ import asyncclick as click
 from anyio import create_task_group
 from pyslobs import StreamingService
 
+from . import console
 from .cli import cli
 from .errors import SlobsCliError
 
@@ -29,7 +30,7 @@ async def start(ctx: click.Context):
             raise SlobsCliError('Replay buffer is already active.')
 
         await ss.start_replay_buffer()
-        click.echo('Replay buffer started.')
+        console.out.print('Replay buffer started.')
         conn.close()
 
     try:
@@ -57,7 +58,7 @@ async def stop(ctx: click.Context):
             raise SlobsCliError('Replay buffer is already inactive.')
 
         await ss.stop_replay_buffer()
-        click.echo('Replay buffer stopped.')
+        console.out.print('Replay buffer stopped.')
         conn.close()
 
     try:
@@ -80,9 +81,9 @@ async def status(ctx: click.Context):
         model = await ss.get_model()
         active = model.replay_buffer_status != 'offline'
         if active:
-            click.echo('Replay buffer is currently active.')
+            console.out.print('Replay buffer is currently active.')
         else:
-            click.echo('Replay buffer is currently inactive.')
+            console.out.print('Replay buffer is currently inactive.')
         conn.close()
 
     async with create_task_group() as tg:
@@ -99,7 +100,7 @@ async def save(ctx: click.Context):
 
     async def _run():
         await ss.save_replay()
-        click.echo('Replay buffer saved.')
+        console.out.print('Replay buffer saved.')
         conn.close()
 
     async with create_task_group() as tg:
diff --git a/src/slobs_cli/scene.py b/src/slobs_cli/scene.py
index 846209a..e12f92b 100644
--- a/src/slobs_cli/scene.py
+++ b/src/slobs_cli/scene.py
@@ -3,8 +3,10 @@
 import asyncclick as click
 from anyio import create_task_group
 from pyslobs import ProtocolError, ScenesService, TransitionsService
-from terminaltables3 import AsciiTable
+from rich.table import Table
+from rich.text import Text
 
+from . import console, util
 from .cli import cli
 from .errors import SlobsCliError, SlobsCliProtocolError
 
@@ -25,34 +27,45 @@ async def list(ctx: click.Context, id: bool = False):
     async def _run():
         scenes = await ss.get_scenes()
         if not scenes:
-            click.echo('No scenes found.')
+            console.out.print('No scenes found.')
             conn.close()
             return
 
         active_scene = await ss.active_scene()
 
-        table_data = [
-            ['Scene Name', 'ID', 'Active'] if id else ['Scene Name', 'Active']
-        ]
+        style = ctx.obj['style']
+        table = Table(
+            show_header=True,
+            header_style=style.header,
+            border_style=style.border,
+        )
+
+        if id:
+            columns = [
+                ('Scene Name', 'left'),
+                ('Active', 'center'),
+                ('ID', 'left'),
+            ]
+        else:
+            columns = [
+                ('Scene Name', 'left'),
+                ('Active', 'center'),
+            ]
+
+        for col_name, col_justify in columns:
+            table.add_column(Text(col_name, justify='center'), justify=col_justify)
+
         for scene in scenes:
-            if scene.id == active_scene.id:
-                to_append = [click.style(scene.name, fg='green')]
-            else:
-                to_append = [click.style(scene.name, fg='blue')]
+            to_append = [Text(scene.name, style=style.cell)]
+            to_append.append(
+                util.check_mark(ctx, scene.id == active_scene.id, empty_if_false=True)
+            )
             if id:
-                to_append.append(scene.id)
-            if scene.id == active_scene.id:
-                to_append.append('✅')
+                to_append.append(Text(scene.id, style=style.cell))
 
-            table_data.append(to_append)
+            table.add_row(*to_append)
 
-        table = AsciiTable(table_data)
-        table.justify_columns = {
-            0: 'left',
-            1: 'left' if id else 'center',
-            2: 'center' if id else None,
-        }
-        click.echo(table.table)
+        console.out.print(table)
 
         conn.close()
 
@@ -76,9 +89,9 @@ async def current(ctx: click.Context, id: bool = False):
 
     async def _run():
         active_scene = await ss.active_scene()
-        click.echo(
-            f'Current active scene: {click.style(active_scene.name, fg="green")} '
-            f'{f"(ID: {active_scene.id})" if id else ""}'
+        console.out.print(
+            f'Current active scene: {console.highlight(ctx, active_scene.name)} '
+            f'{f"(ID: {console.highlight(ctx, active_scene.id)})" if id else ""}'
         )
         conn.close()
 
@@ -118,18 +131,21 @@ async def switch(
                 if model.studio_mode:
                     await ss.make_scene_active(scene.id)
                     if preview:
-                        click.echo(
-                            f'Switched to preview scene: {click.style(scene.name, fg="blue")} '
-                            f'{f"(ID: {scene.id})." if id else ""}'
+                        console.out.print(
+                            f'Switched to preview scene: {console.highlight(ctx, scene.name)} '
+                            f'{f"(ID: {console.highlight(ctx, scene.id)})" if id else ""}'
                         )
                     else:
-                        click.echo(
-                            f'Switched to scene: {click.style(scene.name, fg="blue")} '
-                            f'{f"(ID: {scene.id})." if id else ""}'
+                        console.out.print(
+                            f'Switched to scene: {console.highlight(ctx, scene.name)} '
+                            f'{f"(ID: {console.highlight(ctx, scene.id)})" if id else ""}'
                         )
-                        await ts.execute_studio_mode_transition()
-                        click.echo(
-                            'Executed studio mode transition to make the scene active.'
+                        console.err.print(
+                            console.warning(
+                                ctx,
+                                'Warning: You are in studio mode. The scene switch is not active yet.\n'
+                                'use `slobs-cli studiomode force-transition` to activate the scene switch.',
+                            )
                         )
                 else:
                     if preview:
@@ -139,9 +155,9 @@ async def switch(
                         )
 
                     await ss.make_scene_active(scene.id)
-                    click.echo(
-                        f'Switched to scene: {click.style(scene.name, fg="blue")} '
-                        f'{f"(ID: {scene.id})." if id else ""}'
+                    console.out.print(
+                        f'Switched to scene: {console.highlight(ctx, scene.name)} '
+                        f'{f"(ID: {console.highlight(ctx, scene.id)})" if id else ""}'
                     )
 
                 conn.close()
diff --git a/src/slobs_cli/scenecollection.py b/src/slobs_cli/scenecollection.py
index 3c2322a..254ca0e 100644
--- a/src/slobs_cli/scenecollection.py
+++ b/src/slobs_cli/scenecollection.py
@@ -3,8 +3,10 @@
 import asyncclick as click
 from anyio import create_task_group
 from pyslobs import ISceneCollectionCreateOptions, SceneCollectionsService
-from terminaltables3 import AsciiTable
+from rich.table import Table
+from rich.text import Text
 
+from . import console, util
 from .cli import cli
 from .errors import SlobsCliError
 
@@ -25,35 +27,46 @@ async def list(ctx: click.Context, id: bool):
     async def _run():
         collections = await scs.collections()
         if not collections:
-            click.echo('No scene collections found.')
+            console.out.print('No scene collections found.')
             conn.close()
             return
 
         active_collection = await scs.active_collection()
 
-        table_data = [
-            ['Scene Collection Name', 'ID', 'Active']
-            if id
-            else ['Scene Collection Name', 'Active']
-        ]
-        for collection in collections:
-            if collection.id == active_collection.id:
-                to_append = [click.style(collection.name, fg='green')]
-            else:
-                to_append = [click.style(collection.name, fg='blue')]
-            if id:
-                to_append.append(collection.id)
-            if collection.id == active_collection.id:
-                to_append.append('✅')
-            table_data.append(to_append)
+        style = ctx.obj['style']
+        table = Table(
+            show_header=True,
+            header_style=style.header,
+            border_style=style.border,
+        )
 
-        table = AsciiTable(table_data)
-        table.justify_columns = {
-            0: 'left',
-            1: 'left' if id else 'center',
-            2: 'center' if id else None,
-        }
-        click.echo(table.table)
+        if id:
+            columns = [
+                ('Scene Collection Name', 'left'),
+                ('Active', 'center'),
+                ('ID', 'left'),
+            ]
+        else:
+            columns = [
+                ('Scene Collection Name', 'left'),
+                ('Active', 'center'),
+            ]
+
+        for col_name, col_justify in columns:
+            table.add_column(Text(col_name, justify='center'), justify=col_justify)
+
+        for collection in collections:
+            to_append = [Text(collection.name, style=style.cell)]
+            to_append.append(
+                util.check_mark(
+                    ctx, collection.id == active_collection.id, empty_if_false=True
+                )
+            )
+            if id:
+                to_append.append(Text(collection.id, style=style.cell))
+            table.add_row(*to_append)
+
+        console.out.print(table)
 
         conn.close()
 
@@ -80,7 +93,9 @@ async def load(ctx: click.Context, scenecollection_name: str):
             raise SlobsCliError(f'Scene collection "{scenecollection_name}" not found.')
 
         await scs.load(collection.id)
-        click.echo(f'Scene collection "{scenecollection_name}" loaded successfully.')
+        console.out.print(
+            f'Scene collection {console.highlight(scenecollection_name)} loaded successfully.'
+        )
         conn.close()
 
     try:
@@ -102,7 +117,9 @@ async def create(ctx: click.Context, scenecollection_name: str):
 
     async def _run():
         await scs.create(ISceneCollectionCreateOptions(scenecollection_name))
-        click.echo(f'Scene collection "{scenecollection_name}" created successfully.')
+        console.out.print(
+            f'Scene collection {console.highlight(scenecollection_name)} created successfully.'
+        )
         conn.close()
 
     async with create_task_group() as tg:
@@ -128,7 +145,9 @@ async def delete(ctx: click.Context, scenecollection_name: str):
             raise SlobsCliError(f'Scene collection "{scenecollection_name}" not found.')
 
         await scs.delete(collection.id)
-        click.echo(f'Scene collection "{scenecollection_name}" deleted successfully.')
+        console.out.print(
+            f'Scene collection {console.highlight(scenecollection_name)} deleted successfully.'
+        )
         conn.close()
 
     try:
@@ -159,8 +178,8 @@ async def rename(ctx: click.Context, scenecollection_name: str, new_name: str):
             raise SlobsCliError(f'Scene collection "{scenecollection_name}" not found.')
 
         await scs.rename(new_name, collection.id)
-        click.echo(
-            f'Scene collection "{scenecollection_name}" renamed to "{new_name}".'
+        console.out.print(
+            f'Scene collection {console.highlight(scenecollection_name)} renamed to {console.highlight(new_name)}.'
         )
         conn.close()
 
diff --git a/src/slobs_cli/stream.py b/src/slobs_cli/stream.py
index 3003ac5..d9f1380 100644
--- a/src/slobs_cli/stream.py
+++ b/src/slobs_cli/stream.py
@@ -4,6 +4,7 @@ import asyncclick as click
 from anyio import create_task_group
 from pyslobs import StreamingService
 
+from . import console
 from .cli import cli
 from .errors import SlobsCliError
 
@@ -29,7 +30,7 @@ async def start(ctx: click.Context):
             raise SlobsCliError('Stream is already active.')
 
         await ss.toggle_streaming()
-        click.echo('Stream started.')
+        console.out.print('Stream started.')
         conn.close()
 
     try:
@@ -57,7 +58,7 @@ async def stop(ctx: click.Context):
             raise SlobsCliError('Stream is already inactive.')
 
         await ss.toggle_streaming()
-        click.echo('Stream stopped.')
+        console.out.print('Stream stopped.')
         conn.close()
 
     try:
@@ -81,9 +82,9 @@ async def status(ctx: click.Context):
         active = model.streaming_status != 'offline'
 
         if active:
-            click.echo('Stream is currently active.')
+            console.out.print('Stream is currently active.')
         else:
-            click.echo('Stream is currently inactive.')
+            console.out.print('Stream is currently inactive.')
         conn.close()
 
     async with create_task_group() as tg:
@@ -104,9 +105,9 @@ async def toggle(ctx: click.Context):
 
         await ss.toggle_streaming()
         if active:
-            click.echo('Stream stopped.')
+            console.out.print('Stream stopped.')
         else:
-            click.echo('Stream started.')
+            console.out.print('Stream started.')
 
         conn.close()
 
diff --git a/src/slobs_cli/studiomode.py b/src/slobs_cli/studiomode.py
index 2581145..65e9ed7 100644
--- a/src/slobs_cli/studiomode.py
+++ b/src/slobs_cli/studiomode.py
@@ -4,6 +4,7 @@ import asyncclick as click
 from anyio import create_task_group
 from pyslobs import TransitionsService
 
+from . import console
 from .cli import cli
 from .errors import SlobsCliError
 
@@ -27,7 +28,7 @@ async def enable(ctx: click.Context):
             raise SlobsCliError('Studio mode is already enabled.')
 
         await ts.enable_studio_mode()
-        click.echo('Studio mode enabled successfully.')
+        console.out.print('Studio mode enabled successfully.')
         conn.close()
 
     try:
@@ -53,7 +54,7 @@ async def disable(ctx: click.Context):
             raise SlobsCliError('Studio mode is already disabled.')
 
         await ts.disable_studio_mode()
-        click.echo('Studio mode disabled successfully.')
+        console.out.print('Studio mode disabled successfully.')
         conn.close()
 
     try:
@@ -75,9 +76,9 @@ async def status(ctx: click.Context):
     async def _run():
         model = await ts.get_model()
         if model.studio_mode:
-            click.echo('Studio mode is currently enabled.')
+            console.out.print('Studio mode is currently enabled.')
         else:
-            click.echo('Studio mode is currently disabled.')
+            console.out.print('Studio mode is currently disabled.')
         conn.close()
 
     async with create_task_group() as tg:
@@ -96,10 +97,10 @@ async def toggle(ctx: click.Context):
         model = await ts.get_model()
         if model.studio_mode:
             await ts.disable_studio_mode()
-            click.echo('Studio mode disabled successfully.')
+            console.out.print('Studio mode disabled successfully.')
         else:
             await ts.enable_studio_mode()
-            click.echo('Studio mode enabled successfully.')
+            console.out.print('Studio mode enabled successfully.')
         conn.close()
 
     async with create_task_group() as tg:
@@ -121,7 +122,7 @@ async def force_transition(ctx: click.Context):
             raise SlobsCliError('Studio mode is not enabled.')
 
         await ts.execute_studio_mode_transition()
-        click.echo('Forced studio mode transition.')
+        console.out.print('Forced studio mode transition.')
         conn.close()
 
     try:
diff --git a/src/slobs_cli/util.py b/src/slobs_cli/util.py
new file mode 100644
index 0000000..aacc0d7
--- /dev/null
+++ b/src/slobs_cli/util.py
@@ -0,0 +1,15 @@
+"""module containing utility functions for Slobs CLI."""
+
+import os
+
+import asyncclick as click
+
+
+def check_mark(ctx: click.Context, value: bool, empty_if_false: bool = False) -> str:
+    """Return a check mark or cross mark based on the boolean value."""
+    if empty_if_false and not value:
+        return ''
+
+    if os.getenv('NO_COLOR', '') != '' or ctx.obj['style'].name == 'no_colour':
+        return '✓' if value else '✗'
+    return '✅' if value else '❌'