Coverage for python/lsst/images/cli/_diagram.py: 100%
44 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-04 10:14 -0700
1# This file is part of lsst-images.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
11from __future__ import annotations
13__all__ = ("diagram",)
15import importlib.metadata
17import click
19from ..diagram import build_graph, graph_from_file, make_policy, render
20from ..serialization._asdf_utils import ArrayReferenceModel
21from ..serialization._io import (
22 _BUILTIN_SCHEMA_PROVIDERS,
23 _REGISTRY,
24 _SCHEMA_ENTRY_POINT_GROUP,
25 class_for_schema,
26 parameterize_tree,
27)
30def _available_schemas() -> list[str]:
31 """Return the sorted schema names that can be diagrammed.
33 Importing this module imports ``lsst.images`` first, so every
34 unconditionally-imported model is already registered in ``_REGISTRY``. The
35 lazily-loaded built-in providers and the third-party
36 ``lsst.images.schemas`` entry points are added by name without importing
37 them, so a schema appears here even though `class_for_schema` only loads
38 its provider on demand.
39 """
40 entry_point_names = {
41 entry_point.name for entry_point in importlib.metadata.entry_points(group=_SCHEMA_ENTRY_POINT_GROUP)
42 }
43 return sorted(set(_REGISTRY) | set(_BUILTIN_SCHEMA_PROVIDERS) | entry_point_names)
46@click.command(name="diagram")
47@click.argument("model", required=False)
48@click.option(
49 "--from-file",
50 "from_file",
51 type=click.Path(exists=True, dir_okay=False),
52 help="Diagram the concrete structure of a serialized file instead of a model.",
53)
54@click.option(
55 "--format",
56 "fmt",
57 type=click.Choice(["mermaid", "dot", "tree"]),
58 default="mermaid",
59 show_default=True,
60 help="Output format.",
61)
62@click.option(
63 "--expand",
64 multiple=True,
65 metavar="TYPE",
66 help="Expand a type that is collapsed by default. Repeatable."
67 "The name can match the public name ('Image') or the serialization name "
68 "('ImageSerializationModel').",
69)
70@click.option(
71 "--collapse",
72 multiple=True,
73 metavar="TYPE",
74 help="Render TYPE as a leaf without expanding its fields. "
75 "Repeatable; useful for noisy sub-trees such as --collapse ButlerInfo. "
76 "The name can match the public name ('Image') or the serialization name "
77 "('ImageSerializationModel').",
78)
79@click.option(
80 "--expand-leaves",
81 is_flag=True,
82 help="Expand the serialization-helper leaves that are collapsed by default.",
83)
84@click.option(
85 "--attributes",
86 is_flag=True,
87 help="Show scalar fields too; by default only model composition is shown.",
88)
89@click.option(
90 "--hide-field",
91 "hide_fields",
92 multiple=True,
93 metavar="NAME",
94 help="Drop every field with this name, along with any sub-tree reachable only through it. "
95 "Repeatable; useful for clipping pervasive payload fields such as "
96 "'--hide-field data' and '--hide-field table'.",
97)
98@click.option(
99 "--hide-type",
100 "hide_types",
101 multiple=True,
102 metavar="TYPE",
103 help="Drop a type entirely, removing every edge that points at it (unlike --collapse, which keeps the "
104 "node as a leaf). Repeatable; matches the public or serialization name, e.g. '--hide-type TableModel'.",
105)
106@click.option(
107 "--serialization-names",
108 is_flag=True,
109 help="Label nodes with serialization-model class names instead of public class names.",
110)
111@click.option(
112 "-o", "--output", type=click.Path(dir_okay=False), default=None, help="Write to a file instead of stdout."
113)
114@click.option("--list", "list_models", is_flag=True, help="List available schema names and exit.")
115def diagram(
116 model: str | None,
117 from_file: str | None,
118 fmt: str,
119 expand: tuple[str, ...],
120 collapse: tuple[str, ...],
121 expand_leaves: bool,
122 attributes: bool,
123 hide_fields: tuple[str, ...],
124 hide_types: tuple[str, ...],
125 serialization_names: bool,
126 output: str | None,
127 list_models: bool,
128) -> None: # numpydoc ignore=PR01
129 """Generate a composition diagram of an lsst.images model.
131 Pass a schema name (e.g. ``visit-image`` or ``cell-coadd``) to diagram the
132 abstract model, or ``--from-file`` to diagram the concrete structure of a
133 serialized file, which collapses unions such as the PSF to the type
134 actually stored. Use ``--list`` to see the available schema names.
135 """
136 if list_models:
137 for name in _available_schemas():
138 click.echo(name)
139 return
140 if bool(model) == bool(from_file):
141 raise click.UsageError("Provide exactly one of MODEL or --from-file.")
143 policy = make_policy(
144 expand_leaves=expand_leaves,
145 expand=expand,
146 collapse=collapse,
147 include_attributes=attributes,
148 hide_fields=hide_fields,
149 hide_types=hide_types,
150 public_names=not serialization_names,
151 )
152 if from_file is not None:
153 graph = graph_from_file(from_file, policy=policy)
154 else:
155 assert model is not None
156 tree_cls = class_for_schema(model.replace("-", "_"))
157 if tree_cls is None:
158 available = ", ".join(_available_schemas())
159 raise click.ClickException(f"Unknown model {model!r}; choose from: {available}.")
160 graph = build_graph(parameterize_tree(tree_cls, ArrayReferenceModel), policy=policy)
162 text = render(graph, fmt)
163 if output is not None:
164 with open(output, "w") as stream:
165 stream.write(text)
166 else:
167 click.echo(text, nl=False)