Fix installing Jockey as a dependency and introduce AZ filter

pyproject.toml

@@ -1,14 +1,13 @@
 [project]
 requires-python = ">=3.10"
 
-[tool.poetry]
 name = "jockey"
 version = "0.2.1"
 description = "A CLI tool designed to facilitate quick and easy retrieval of Juju objects using filters."
 readme = "README.md"
 authors = [
-    "Connor Chamberlain <connor.chamberlain@canonical.com>",
-    "Pedro Castillo <pedro.castillo@canonical.com>"
+    { name = "Connor Chamberlain", email = "connor.chamberlain@canonical.com" },
+    { name = "Pedro Castillo", email = "pedro.castillo@canonical.com" },
 ]
 
 classifiers = [
@@ -17,12 +16,19 @@ classifiers = [
     "Operating System :: POSIX :: Linux",
 ]
 
-packages = [
-    { include = "jockey", from = "src" }
+[tool.poetry]
+name = "jockey"
+version = "0.2.1"
+description = "A CLI tool designed to facilitate quick and easy retrieval of Juju objects using filters."
+readme = "README.md"
+authors = [
+    "Connor Chamberlain <connor.chamberlain@canonical.com>",
+    "Pedro Castillo <pedro.castillo@canonical.com>",
 ]
+packages = [{ include = "jockey", from = "src" }]
 
 [tool.poetry.scripts]
-juju-jockey = "jockey.__main__:main"
+juju-jockey = "jockey:main"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<3.13"
@@ -72,9 +78,7 @@ directory = "htmlcov"
 
 [tool.pytest.ini_options]
 pythonpath = "src"
-addopts = [
-    "--import-mode=importlib",
-]
+addopts = ["--import-mode=importlib"]
 md_report = true
 md_report_verbose = 0
 md_report_color = "auto"
@@ -145,14 +149,11 @@ version = "6.7.0"
 disable-upx = true
 
 [tool.poetry-pyinstaller-plugin.scripts]
-juju-jockey = { source = "src/jockey/__main__.py", type = "onefile", bundle = false }
+juju-jockey = { source = "src/jockey/__init__.py", type = "onefile", bundle = false }
 
 [tool.commitizen]
 tag_format = "v$version"
-version_files = [
-    "pyproject.toml:version",
-    "src/jockey/__init__.py"
-]
+version_files = ["pyproject.toml:version", "src/jockey/__init__.py"]
 bump_message = "bump: $current_version -> $new_version [skip ci]"
 
 [build-system]

src/jockey/__init__.py

@@ -13,3 +13,62 @@
 __copyright__ = "Copyright (c) 2024, Connor Chamberlain"
 
 __docformat__ = "numpy"
+
+import logging
+import os
+from pkgutil import get_data
+import sys
+from typing import Optional, Sequence
+
+from rich import print
+from rich.console import Console
+from rich.markdown import Markdown
+
+from jockey.__args__ import parse_args
+from jockey.core import query
+from jockey.log import configure_logging
+
+
+logger = logging.getLogger(__name__)
+
+
+if "SNAP" in os.environ:
+    os.environ["PATH"] += ":" + os.path.join(os.environ["SNAP"], "usr", "juju", "bin")
+
+
+def info() -> Markdown:
+    info_data = get_data("jockey", "info.md")
+    info_decoded = info_data.decode("utf-8") if info_data else ""
+    return Markdown(info_decoded)
+
+
+def print_info(console: Optional[Console] = None) -> None:
+    if not console:
+        console = Console(width=120)
+
+    console.print(info())
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    # parse command-line arguments and configure logging
+    if argv is None:
+        argv = sys.argv[1:]
+
+    args = parse_args(argv)
+    logger.debug("Parsed command-line arguments:\n%r", args)
+
+    verbosity = args.verbose if "verbose" in args else 0
+    configure_logging(verbosity)
+
+    # check if OBJECT is requesting the informational message
+    if args.object == "info":
+        print_info()
+        return 0
+
+    filters = args.filters if "filters" in args else []
+    print("\n".join(query(object_type=args.object, filter_strings=filters, file=args.file, model=args.model)))
+    return 0
+
+
+if __name__ == "__main__":
+    exit(main(sys.argv[1:]))

src/jockey/core.py

@@ -48,6 +48,7 @@ class ObjectType(Enum):
     MACHINE = ("machines", "machine", "m")
     IP = ("ips", "address", "addresses", "ip", "i")
     HOSTNAME = ("hostnames", "hostname", "host", "hosts", "h")
+    AVAILABILITY_ZONE = ("availability-zone", "availability_zone", "az", "zone")
 
 
 def list_abbreviations() -> str:
@@ -683,6 +684,33 @@ def ip_to_machine(status: JujuStatus, ip: str) -> str:
     raise Exception(f"No machine found with IP {ip}")
 
 
+def machine_to_availability_zone(status: JujuStatus, machine: str) -> Optional[str]:
+    """
+    Given an machine id, get its availability zone.
+
+    Arguments
+    =========
+    status (JujuStatus)
+        The current Juju status in json format.
+    machine (str)
+        The ID of the machine to use.
+
+    Returns
+    =======
+    avaliability_zone (str)
+        The machine's avilability zone.
+    """
+    if "lxd" in machine:
+        # A container's availability zone is the same as its parent, so we grab that.
+        machine, _, _ = machine.split("/")
+    hardware = status["machines"][machine]["hardware"]
+    for entry in hardware.split():
+        key, value = entry.split("=")
+        if key == "availability-zone":
+            return value
+    return None
+
+
 def machine_to_hostname(status: JujuStatus, machine: str) -> str:
     """
     Given an machine id, get its hostname.
@@ -751,6 +779,7 @@ def filter_units(status: JujuStatus, filters: List[JockeyFilter]) -> Generator[s
     machine_filters = [f for f in filters if f.obj_type == ObjectType.MACHINE]
     ip_filters = [f for f in filters if f.obj_type == ObjectType.IP]
     hostname_filters = [f for f in filters if f.obj_type == ObjectType.HOSTNAME]
+    az_filters = [f for f in filters if f.obj_type == ObjectType.AVAILABILITY_ZONE]
 
     for unit in get_units(status):
         # Check unit filters
@@ -793,6 +822,11 @@ def filter_units(status: JujuStatus, filters: List[JockeyFilter]) -> Generator[s
         if not all(any(check_filter_match(i_filter, ip) for ip in ips) for i_filter in ip_filters):
             continue
 
+        # Check availability zone filter
+        az = machine_to_availability_zone(status, machine)
+        if not all(check_filter_match(az_filter, az) for az_filter in az_filters):
+            continue
+
         yield unit
 
 
@@ -816,6 +850,7 @@ def filter_machines(status: JujuStatus, filters: List[JockeyFilter]) -> Generato
     machine_filters = [f for f in filters if f.obj_type == ObjectType.MACHINE]
     hostname_filters = [f for f in filters if f.obj_type == ObjectType.HOSTNAME]
     ip_filters = [f for f in filters if f.obj_type == ObjectType.IP]
+    az_filters = [f for f in filters if f.obj_type == ObjectType.AVAILABILITY_ZONE]
 
     unit_filters = [f for f in filters if f.obj_type == ObjectType.UNIT]
     app_filters = [f for f in filters if f.obj_type == ObjectType.APP]
@@ -854,6 +889,11 @@ def filter_machines(status: JujuStatus, filters: List[JockeyFilter]) -> Generato
         if not check_filter_batch_match(charm_filters, charms):  # type: ignore
             continue
 
+        # Check availability zone filter
+        az = machine_to_availability_zone(status, machine)
+        if not all(check_filter_match(az_filter, az) for az_filter in az_filters):
+            continue
+
         yield machine
 
 

tests/test_cli.py

@@ -6,7 +6,7 @@
 
 import pytest
 
-from jockey.__main__ import main
+from jockey import main
 from tests.test_util import SAMPLES_DIR, StandardOutputCapture