feat(operations.pkg): migrate FreeBSD pkg to PackageInfo (phase 2 of #1725)

src/pyinfra/facts/pkg.py

@@ -1,30 +1,86 @@
 from __future__ import annotations
 
+import re
+
 from typing_extensions import override
 
 from pyinfra.api import FactBase
 
-from .util.packaging import parse_packages
+from .util.packages import PackageInfo, build_package_map
+
+# Separates installed-package lines from the `pkg version` remote comparison
+# lines in the compound command output below.
+_REMOTE_MARKER = "__pyinfra_pkg_remote__"
 
+_REMOTE_HAS_RE = re.compile(r"\(remote has ([^)]+)\)")
 
-class PkgPackages(FactBase):
+
+class PkgPackages(FactBase[list[PackageInfo]]):
     """
-    Returns a dict of installed pkg packages:
+    Returns a list of installed pkg packages as ``PackageInfo``, sorted by
+    name. On FreeBSD (``pkg``) entries carry the locked status (``HELD``, via
+    ``pkg lock``) and available upgrades (``UPGRADEABLE``, compared against
+    the remote catalogues). With the ``pkg_info`` fallback used on other BSDs
+    every entry is ``INSTALLED``.
 
     .. code:: python
 
-        {
-            "package_name": ["version"],
-        }
+        [
+            PackageInfo(
+                name="zsh",
+                installed_versions=("5.9_5",),
+                available_version="5.9_6",
+                status=PackageStatus.UPGRADEABLE,
+            ),
+        ]
     """
 
     regex = r"^([a-zA-Z0-9_\-\+]+)\-([0-9a-z\.]+)"
-    default = dict
+    default = list
 
     @override
     def command(self) -> str:
-        return "pkg info || pkg_info || true"
+        return (
+            f"(pkg query '%n %v %k' && echo {_REMOTE_MARKER} && (pkg version -vRL= || true))"
+            " || pkg_info || true"
+        )
 
     @override
-    def process(self, output):
-        return parse_packages(self.regex, output)
+    def process(self, output: list[str]) -> list[PackageInfo]:
+        installed: dict[str, set[str]] = {}
+        upgradeable: dict[str, str] = {}
+        held: set[str] = set()
+
+        # The marker only appears when `pkg query` ran, so its presence tells
+        # us whether to parse `name version locked` fields or pkg_info lines.
+        is_pkg_query = any(line.strip() == _REMOTE_MARKER for line in output)
+
+        in_remote = False
+        for line in output:
+            if line.strip() == _REMOTE_MARKER:
+                in_remote = True
+                continue
+
+            if in_remote:
+                # `pkg version -vRL=` line: `name-version  <  needs updating (remote has X)`
+                parts = line.split()
+                if len(parts) >= 2 and parts[1] == "<":
+                    match = _REMOTE_HAS_RE.search(line)
+                    if match:
+                        upgradeable[parts[0].rsplit("-", 1)[0]] = match.group(1)
+                continue
+
+            if is_pkg_query:
+                # `pkg query '%n %v %k'` line: `name version locked(0|1)`
+                parts = line.split()
+                if len(parts) == 3 and parts[2] in ("0", "1"):
+                    name, version, locked = parts
+                    installed.setdefault(name, set()).add(version)
+                    if locked == "1":
+                        held.add(name)
+            else:
+                match = re.match(self.regex, line)
+                if match:
+                    installed.setdefault(match.group(1), set()).add(match.group(2))
+
+        return build_package_map(installed, upgradeable, held)

src/pyinfra/facts/util/packages.py

@@ -6,6 +6,8 @@
 ``dict[str, set[str]]``.
 """
 
+from __future__ import annotations
+
 import re
 from dataclasses import dataclass
 from enum import Enum
@@ -57,29 +59,46 @@ class PackageInfo:
     def installed_version(self) -> str | None:
         return self.installed_versions[-1] if self.installed_versions else None
 
+    @classmethod
+    def from_dict(cls, data: dict) -> PackageInfo:
+        """Build a :class:`PackageInfo` from a dict-shaped description.
+
+        This is useful for consumers that receive plain JSON-like package
+        data (for example test fixtures or inventory variables) and need a
+        real :class:`PackageInfo` instance.
+        """
+        return cls(
+            name=data["name"],
+            installed_versions=tuple(data.get("installed_versions", ())),
+            available_version=data.get("available_version"),
+            status=PackageStatus(data.get("status", "installed")),
+        )
+
 
 def build_package_map(
     installed: dict[str, set[str]],
     upgradeable: dict[str, str] | None = None,
     held: set[str] | None = None,
-) -> dict[str, PackageInfo]:
-    """Build a :class:`PackageInfo` map by combining sub-fact data.
+) -> list[PackageInfo]:
+    """Build a list of :class:`PackageInfo` by combining package data.
 
-    + installed: installed packages from a fact (name to set of versions).
+    + installed: installed packages (name to set of versions).
     + upgradeable: packages with available upgrades (name to available version).
     + held: names of held/locked/pinned packages.
 
-    Versions are sorted with a natural-order key (digit runs as integers) so
-    multi-version output is deterministic across runs and the highest version
-    is last.
+    The result is a flat ``list[PackageInfo]`` sorted by name (each entry
+    carries its own name, so a name-keyed dict would only duplicate it).
+    Versions within an entry are sorted with a natural-order key (digit runs
+    as integers) so multi-version output is deterministic across runs and the
+    highest version is last.
     """
 
-    result: dict[str, PackageInfo] = {}
+    result: list[PackageInfo] = []
     _upgradeable = upgradeable or {}
     _held = held or set()
 
-    for name, versions in installed.items():
-        sorted_versions = tuple(sorted(versions, key=_version_sort_key))
+    for name in sorted(installed):
+        sorted_versions = tuple(sorted(installed[name], key=_version_sort_key))
 
         if name in _held:
             status = PackageStatus.HELD
@@ -88,11 +107,13 @@ def build_package_map(
         else:
             status = PackageStatus.INSTALLED
 
-        result[name] = PackageInfo(
-            name=name,
-            installed_versions=sorted_versions,
-            available_version=_upgradeable.get(name),
-            status=status,
+        result.append(
+            PackageInfo(
+                name=name,
+                installed_versions=sorted_versions,
+                available_version=_upgradeable.get(name),
+                status=status,
+            )
         )
 
     return result

src/pyinfra/operations/pkg.py

@@ -9,20 +9,70 @@
 from pyinfra.facts.files import File
 from pyinfra.facts.pkg import PkgPackages
 from pyinfra.facts.server import Arch, Os, OsVersion, Which
+from pyinfra.facts.util.packages import PackageInfo, PackageStatus
 
 from .util.packaging import ensure_packages
 
 
+@operation(is_idempotent=False)
+def update(force: bool = False):
+    """
+    Update the local pkg repository catalogues (FreeBSD ``pkg`` only).
+
+    + force: force a full download of the repository catalogues
+    """
+
+    yield "pkg update -f" if force else "pkg update"
+
+
 @operation()
-def packages(packages: str | list[str] | None = None, present=True, pkg_path: str | None = None):
+def upgrade():
+    """
+    Upgrade all pkg packages with an available upgrade (FreeBSD ``pkg`` only).
+
+    Noops when no installed package has an upgrade available.
+
+    **Example:**
+
+    .. code:: python
+
+        from pyinfra.operations import pkg
+        pkg.upgrade(
+            name="Upgrade all pkg packages",
+        )
+
+    """
+
+    current_packages = [
+        PackageInfo.from_dict(package) if isinstance(package, dict) else package
+        for package in host.get_fact(PkgPackages)
+    ]
+
+    if not any(package.status is PackageStatus.UPGRADEABLE for package in current_packages):
+        host.noop("all packages are up to date")
+        return
+
+    yield "pkg upgrade -y"
+
+
+@operation()
+def packages(
+    packages: str | list[str] | None = None,
+    present: bool = True,
+    latest: bool = False,
+    pkg_path: str | None = None,
+):
     """
     Install/remove/update pkg packages. This will use ``pkg ...`` where available
     (FreeBSD) and the ``pkg_*`` variants elsewhere.
 
     + packages: list of packages to ensure
     + present: whether the packages should be installed
+    + latest: upgrade installed packages with an available upgrade (FreeBSD ``pkg`` only)
     + pkg_path: the PKG_PATH environment variable to set
 
+    Packages locked with ``pkg lock`` always noop, even with ``latest=True``.
+
     pkg_path:
         By default this is autogenerated as follows (tested/working for OpenBSD):
         ``http://ftp.<OS>.org/pub/<OS>/<VERSION>/packages/<ARCH>/``. Note that OpenBSD's
@@ -56,11 +106,18 @@ def packages(packages: str | list[str] | None = None, present=True, pkg_path: st
     if pkg_path:
         install_command = f"PKG_PATH={pkg_path} {install_command}"
 
+    current_packages = [
+        PackageInfo.from_dict(package) if isinstance(package, dict) else package
+        for package in host.get_fact(PkgPackages)
+    ]
+
     yield from ensure_packages(
         host,
         packages,
-        host.get_fact(PkgPackages),
+        current_packages,
         present,
         install_command=install_command,
         uninstall_command=uninstall_command,
+        latest=latest,
+        upgrade_command="pkg upgrade -y" if is_pkg else None,
     )

src/pyinfra/operations/util/packaging.py

@@ -176,7 +176,7 @@ def _format_version(
 def ensure_packages(
     host: Host,
     packages_to_ensure: str | list[str] | list[PkgInfo] | None,
-    current_packages: dict[str, set[str]] | dict[str, PackageInfo],
+    current_packages: dict[str, set[str]] | dict[str, PackageInfo] | list[PackageInfo],
     present: bool,
     install_command: str | StringCommand,
     uninstall_command: str | StringCommand,
@@ -190,12 +190,13 @@ def ensure_packages(
     Handles this common scenario:
 
     + We have a list of packages(/versions/urls) to ensure
-    + We have a map of existing package -> versions (old) or PackageInfo (new)
+    + We have the existing packages: ``dict[str, set[str]]`` (old),
+      ``dict[str, PackageInfo]`` or ``list[PackageInfo]`` (new)
     + We have the common command bits (install, uninstall, version "joiner")
     + Outputs commands to ensure our desired packages/versions
     + Optionally upgrades packages w/o specified version when present
 
-    When ``current_packages`` values are :class:`PackageInfo` objects, the richer
+    When ``current_packages`` carries :class:`PackageInfo` objects, the richer
     status information is used:
 
     * **HELD** packages always produce a noop, even when ``latest=True``.
@@ -207,7 +208,8 @@ def ensure_packages(
 
     Args:
         packages_to_ensure (list): list of packages or package/versions or PkgInfo's
-        current_packages (dict): dict of package names -> version, or name -> PackageInfo
+        current_packages: dict of package names -> versions, dict of name -> \
+            PackageInfo, or list of PackageInfo
         present (bool): whether packages should exist or not
         install_command (str): command to prefix to list of packages to install
         uninstall_command (str): as above for uninstalling packages
@@ -226,6 +228,9 @@ def ensure_packages(
     if len(packages_to_ensure) == 0:
         return
 
+    if isinstance(current_packages, list):
+        current_packages = {package.name: package for package in current_packages}
+
     packages: list[PkgInfo] = []
     if isinstance(packages_to_ensure[0], PkgInfo):
         packages = cast("list[PkgInfo]", packages_to_ensure)

src/pyinfra_cli/util.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import dataclasses
+import enum
 import json
 import os
 from datetime import datetime
@@ -127,6 +129,12 @@ def json_encode(obj):
     if hasattr(obj, "to_json"):
         return obj.to_json()
 
+    if isinstance(obj, enum.Enum):
+        return obj.value
+
+    if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
+        return dataclasses.asdict(obj)
+
     raise TypeError(f"Cannot serialize: {type(obj)} ({obj})")
 
 

tests/facts/pkg.PkgPackages/empty.yaml

@@ -0,0 +1,3 @@
+command: (pkg query '%n %v %k' && echo __pyinfra_pkg_remote__ && (pkg version -vRL= || true)) || pkg_info || true
+output: []
+fact: []

tests/facts/pkg.PkgPackages/fallback_pkg_info.yaml

@@ -0,0 +1,13 @@
+command: (pkg query '%n %v %k' && echo __pyinfra_pkg_remote__ && (pkg version -vRL= || true)) || pkg_info || true
+output: |
+  nano-3.48 Nano editor
+  py-pip-20.1.1 Python pip
+fact:
+  - name: nano
+    installed_versions: ["3.48"]
+    available_version: null
+    status: installed
+  - name: py-pip
+    installed_versions: ["20.1.1"]
+    available_version: null
+    status: installed

tests/facts/pkg.PkgPackages/freebsd.yaml

@@ -0,0 +1,29 @@
+command: (pkg query '%n %v %k' && echo __pyinfra_pkg_remote__ && (pkg version -vRL= || true)) || pkg_info || true
+output: |
+  zsh 5.9_5 0
+  curl 8.7.1 0
+  vim 9.1.0 1
+  compat-lib 1.0 0
+  compat-lib 1.0_1 0
+  __pyinfra_pkg_remote__
+  zsh-5.9_5                          <   needs updating (remote has 5.9_6)
+  vim-9.1.0                          <   needs updating (remote has 9.1.1)
+  weird-1.0                          >   succeeds remote (remote has 0.9)
+  orphan-2.0                         ?   orphaned: misc/orphan
+fact:
+  - name: compat-lib
+    installed_versions: ["1.0", "1.0_1"]
+    available_version: null
+    status: installed
+  - name: curl
+    installed_versions: ["8.7.1"]
+    available_version: null
+    status: installed
+  - name: vim
+    installed_versions: ["9.1.0"]
+    available_version: "9.1.1"
+    status: held
+  - name: zsh
+    installed_versions: ["5.9_5"]
+    available_version: "5.9_6"
+    status: upgradeable