From d805bfa6ec9dbc25e921fe834beaba57d3dff83b Mon Sep 17 00:00:00 2001 From: John | Elite Encoder Date: Fri, 14 Nov 2025 12:40:20 -0500 Subject: [PATCH] feat: Implement CLI and environment port overrides for StreamProcessor - Added runtime_args module to handle port overrides via CLI and environment variables. - Updated example scripts to utilize the new port resolution logic. - Modified StreamProcessor to support dynamic port assignment based on overrides. - Enhanced logging in examples to reflect the effective port in use. - Introduced tests for runtime_args to ensure correct behavior of port resolution. --- .vscode/launch.json | 4 + examples/__init__.py | 4 + examples/grayscale_chipmunk_example.py | 7 +- examples/passthrough_example.py | 13 +- examples/process_video_example.py | 8 +- pytrickle/runtime_args.py | 172 +++++++++++++++++++++++++ pytrickle/stream_processor.py | 9 +- tests/test_runtime_args.py | 76 +++++++++++ 8 files changed, 278 insertions(+), 15 deletions(-) create mode 100644 pytrickle/runtime_args.py create mode 100644 tests/test_runtime_args.py diff --git a/.vscode/launch.json b/.vscode/launch.json index b3af0ae..0d0a9ca 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -11,6 +11,10 @@ "program": "${workspaceFolder}/examples/process_video_example.py", "console": "integratedTerminal", "cwd": "${workspaceFolder}/examples", + "args": [ + "--port", + "8005" + ], "env": { "ORCH_URL": "https://localhost:9995", "ORCH_SECRET": "orch-secret", diff --git a/examples/__init__.py b/examples/__init__.py index b172d1a..ee8d144 100644 --- a/examples/__init__.py +++ b/examples/__init__.py @@ -3,3 +3,7 @@ This package contains example implementations and utilities for working with pytrickle. """ + +from pytrickle import runtime_args as _runtime_args + +_runtime_args.enable_cli_port_override() diff --git a/examples/grayscale_chipmunk_example.py b/examples/grayscale_chipmunk_example.py index baeff97..55fda7d 100644 --- a/examples/grayscale_chipmunk_example.py +++ b/examples/grayscale_chipmunk_example.py @@ -18,7 +18,7 @@ import numpy as np import torch -from pytrickle import StreamProcessor, VideoFrame, AudioFrame +from pytrickle import StreamProcessor, VideoFrame, AudioFrame, runtime_args from pytrickle.decorators import ( audio_handler, model_loader, @@ -155,14 +155,15 @@ async def on_stop(self) -> None: async def main() -> None: handlers = GrayscaleChipmunkHandlers() + port = runtime_args.resolve_port() processor = StreamProcessor.from_handlers( handlers, name="grayscale-chipmunk-demo", - port=8000, + port=port, ) # Explicitly call load_model to initialize the handlers - logger.info("Initializing handlers...") + logger.info("Initializing handlers on port %d...", port) await processor._frame_processor.load_model() await processor.run_forever() diff --git a/examples/passthrough_example.py b/examples/passthrough_example.py index 1a4d19e..ea16cde 100644 --- a/examples/passthrough_example.py +++ b/examples/passthrough_example.py @@ -18,7 +18,7 @@ from dataclasses import dataclass from typing import List -from pytrickle import StreamProcessor, VideoFrame, AudioFrame +from pytrickle import StreamProcessor, VideoFrame, AudioFrame, runtime_args from pytrickle.decorators import ( audio_handler, model_loader, @@ -125,15 +125,16 @@ async def on_stop(self) -> None: async def main() -> None: """Main entry point - creates and runs the stream processor.""" handlers = PassthroughHandlers() + port = runtime_args.resolve_port() processor = StreamProcessor.from_handlers( handlers, name="passthrough-example", - port=8000, + port=port, ) - - logger.info("Send video to: http://localhost:8000/stream") - logger.info("Update params: POST http://localhost:8000/control") - + + logger.info("Send video to: http://localhost:%d/stream", port) + logger.info("Update params: POST http://localhost:%d/control", port) + await processor.run_forever() diff --git a/examples/process_video_example.py b/examples/process_video_example.py index ac71e44..efa15e1 100644 --- a/examples/process_video_example.py +++ b/examples/process_video_example.py @@ -23,7 +23,7 @@ import numpy as np import torch -from pytrickle import StreamProcessor, VideoFrame +from pytrickle import StreamProcessor, VideoFrame, runtime_args from pytrickle.decorators import ( model_loader, on_stream_start, @@ -250,10 +250,12 @@ async def update_params(params: dict) -> None: async def main() -> None: """Main entry point - creates and runs the stream processor.""" handlers = GreenProcessorHandlers() + port = runtime_args.resolve_port() + processor = StreamProcessor.from_handlers( handlers, name="green-processor", - port=8000, + port=port, frame_skip_config=FrameSkipConfig(), # Optional frame skipping ) @@ -261,7 +263,7 @@ async def main() -> None: handlers.processor = processor logger.info("OpenCV will apply: horizontal flip + green hue") - logger.info("Update params: POST http://localhost:8000/control") + logger.info("Update params: POST http://localhost:%d/control", port) await processor.run_forever() diff --git a/pytrickle/runtime_args.py b/pytrickle/runtime_args.py new file mode 100644 index 0000000..58d8404 --- /dev/null +++ b/pytrickle/runtime_args.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import os +import sys +from pathlib import Path +from typing import List, Optional, Sequence + +_PORT_FLAG = "--port" +_PORT_ENV_VAR = "PYTRICKLE_PORT" + +_cli_override_enabled = False +_cli_port_override: Optional[int] = None +_env_port_override: Optional[int] = None +_override_source: Optional[str] = None + + +def _coerce_port(value: str) -> int: + """Convert *value* to an integer port, raising on invalid input.""" + try: + port = int(value) + except (TypeError, ValueError) as exc: + raise SystemExit( + f"PyTrickle: invalid port '{value}'. Provide an integer between 1 and 65535." + ) from exc + + if not 0 < port < 65536: + raise SystemExit( + f"PyTrickle: invalid port '{value}'. Provide an integer between 1 and 65535." + ) + return port + + +def _parse_cli_port(args: Sequence[str]) -> Optional[int]: + """ + Extract the last --port occurrence from *args* (values after '--' are ignored). + Returns None if flag not present. + """ + port: Optional[int] = None + i = 0 + args_list: List[str] = list(args) + + while i < len(args_list): + token = args_list[i] + + if token == "--": + break + + if token == _PORT_FLAG: + if i + 1 >= len(args_list): + raise SystemExit("PyTrickle: --port flag requires an integer value.") + port = _coerce_port(args_list[i + 1]) + i += 2 + continue + + if token.startswith(f"{_PORT_FLAG}="): + port = _coerce_port(token.split("=", 1)[1]) + i += 1 + continue + + i += 1 + + return port + + +def _load_env_override() -> Optional[int]: + """Read PYTRICKLE_PORT from the environment if present.""" + value = os.getenv(_PORT_ENV_VAR) + if value is None: + return None + port = _coerce_port(value) + return port + + +_env_port_override = _load_env_override() +if _env_port_override is not None: + _override_source = "env" + + +def enable_cli_port_override(args: Optional[Sequence[str]] = None) -> None: + """ + Enable parsing of the process argv for --port. + + Args: + args: Optional iterable of strings to parse instead of sys.argv[1:]. + """ + global _cli_override_enabled + global _cli_port_override + global _override_source + + if _cli_override_enabled: + return + + _cli_override_enabled = True + + parsed_args = args if args is not None else sys.argv[1:] + port = _parse_cli_port(parsed_args) + if port is not None: + _cli_port_override = port + _override_source = "cli" + +def resolve_port(port: Optional[int] = None) -> int: + """ + Return the effective listening port after applying CLI/env overrides. + + Args: + port: Port requested by caller. + + Raises: + ValueError: if no explicit port is provided and no override is enabled. + """ + _auto_enable_cli_override_for_examples() + + if _cli_override_enabled and _cli_port_override is not None: + return _cli_port_override + + if _env_port_override is not None: + return _env_port_override + + if port is None: + raise ValueError("PyTrickle: port must be provided when no CLI or env override is configured.") + + return port + + + +def port_override_source() -> Optional[str]: + """Return 'cli', 'env', or None depending on how the port is overridden.""" + return _override_source + + +def _auto_enable_cli_override_for_examples() -> None: + """Enable CLI overrides automatically when running bundled examples.""" + if _cli_override_enabled: + return + + main_module = sys.modules.get("__main__") + if main_module is None: + return + + main_path = getattr(main_module, "__file__", None) + if not main_path: + return + + try: + main_file = Path(main_path).resolve() + except OSError: + return + + pkg_dir = Path(__file__).resolve().parent + + candidate_dirs = [ + pkg_dir.parent / "examples", + pkg_dir / "examples", + ] + + for candidate in candidate_dirs: + try: + candidate_path = candidate.resolve() + except OSError: + continue + + if not candidate_path.exists(): + continue + + try: + main_file.relative_to(candidate_path) + except ValueError: + continue + + enable_cli_port_override() + return + diff --git a/pytrickle/stream_processor.py b/pytrickle/stream_processor.py index 3686c39..f601c3a 100644 --- a/pytrickle/stream_processor.py +++ b/pytrickle/stream_processor.py @@ -3,6 +3,7 @@ import logging from typing import Optional, Callable, Dict, Any, List, Union, Awaitable +from . import runtime_args from .registry import HandlerRegistry from .frames import VideoFrame, AudioFrame from .frame_processor import FrameProcessor @@ -107,7 +108,7 @@ def __init__( on_stream_stop: Optional async function called when stream stops send_data_interval: Interval for sending data name: Processor name - port: Server port + port: Server port (may be overridden by runtime CLI/env overrides) frame_skip_config: Optional frame skipping configuration (None = no frame skipping) **server_kwargs: Additional arguments passed to StreamServer """ @@ -136,9 +137,11 @@ def __init__( self.param_updater = param_updater self.on_stream_start = on_stream_start self.on_stream_stop = on_stream_stop + resolved_port = runtime_args.resolve_port(port) + self.send_data_interval = send_data_interval self.name = name - self.port = port + self.port = resolved_port self.frame_skip_config = frame_skip_config self.server_kwargs = server_kwargs self._handler_registry: Optional[HandlerRegistry] = None @@ -157,7 +160,7 @@ def __init__( # Create and start server self.server = StreamServer( frame_processor=self._frame_processor, - port=port, + port=resolved_port, frame_skip_config=frame_skip_config, **server_kwargs ) diff --git a/tests/test_runtime_args.py b/tests/test_runtime_args.py new file mode 100644 index 0000000..ea14703 --- /dev/null +++ b/tests/test_runtime_args.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import importlib +import sys +from pathlib import Path +from typing import Optional, Sequence + +import pytest + +import pytrickle.runtime_args as runtime_args + + +def _reload_runtime( + monkeypatch: pytest.MonkeyPatch, + *, + argv: Sequence[str], + env_value: Optional[str] = None, +) -> object: + """Reload runtime_args with controlled argv/env state.""" + monkeypatch.setattr(sys, "argv", list(argv)) + + if env_value is None: + monkeypatch.delenv("PYTRICKLE_PORT", raising=False) + else: + monkeypatch.setenv("PYTRICKLE_PORT", env_value) + + return importlib.reload(runtime_args) + + +def test_env_override_applies_without_cli(monkeypatch: pytest.MonkeyPatch) -> None: + mod = _reload_runtime(monkeypatch, argv=["prog"], env_value="8123") + + assert mod.resolve_port(8000) == 8123 + assert mod.port_override_source() == "env" + + +def test_cli_override_requires_opt_in(monkeypatch: pytest.MonkeyPatch) -> None: + mod = _reload_runtime(monkeypatch, argv=["prog", "--port", "9005"]) + + assert mod.resolve_port(8000) == 8000 + assert mod.port_override_source() is None + + +def test_cli_override_activates_when_enabled(monkeypatch: pytest.MonkeyPatch) -> None: + mod = _reload_runtime(monkeypatch, argv=["prog", "--foo", "bar", "--port", "9001"]) + + mod.enable_cli_port_override() + + assert mod.resolve_port(8000) == 9001 + assert mod.port_override_source() == "cli" + + +def test_cli_override_supports_equals_notation(monkeypatch: pytest.MonkeyPatch) -> None: + mod = _reload_runtime(monkeypatch, argv=["prog", "--port=9100", "--other"]) + + mod.enable_cli_port_override() + assert mod.resolve_port(8000) == 9100 + + +def test_cli_override_invalid_value(monkeypatch: pytest.MonkeyPatch) -> None: + mod = _reload_runtime(monkeypatch, argv=["prog", "--port", "not-a-number"]) + + with pytest.raises(SystemExit): + mod.enable_cli_port_override() + + +def test_repo_examples_trigger_auto_enable(monkeypatch: pytest.MonkeyPatch) -> None: + repo_root = Path(__file__).resolve().parent.parent + example_path = repo_root / "examples" / "loading_overlay_example.py" + + mod = _reload_runtime(monkeypatch, argv=["prog", "--port", "7777"]) + + monkeypatch.setattr(sys.modules["__main__"], "__file__", str(example_path)) + + assert mod.resolve_port(8000) == 7777 +