Coverage for python/lsst/images/cli/_diagram.py: 100%

44 statements  

« prev     ^ index     » next       coverage.py v7.15.0, created at 2026-07-04 17:17 +0000

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 

12 

13__all__ = ("diagram",) 

14 

15import importlib.metadata 

16 

17import click 

18 

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) 

28 

29 

30def _available_schemas() -> list[str]: 

31 """Return the sorted schema names that can be diagrammed. 

32 

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) 

44 

45 

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. 

130 

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.") 

142 

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) 

161 

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)