|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
| 3 | +import ast |
3 | 4 | import importlib.metadata |
| 5 | +import inspect |
| 6 | +import pkgutil |
4 | 7 | import sys |
5 | 8 | from contextlib import contextmanager |
6 | 9 | from datetime import datetime |
|
10 | 13 |
|
11 | 14 | from typing_extensions import Any |
12 | 15 |
|
| 16 | +import ridgeplot as ridgeplot_pkg |
| 17 | + |
13 | 18 | try: |
14 | 19 | from ridgeplot_examples import ALL_EXAMPLES |
15 | 20 | except ImportError: |
|
271 | 276 |
|
272 | 277 |
|
273 | 278 | # Type aliases |
274 | | -_TYPE_ALIASES_FULLY_QUALIFIED = { |
275 | | - # ------- ._color.css_colors ------------------- |
276 | | - "ridgeplot._color.css_colors.CssNamedColor", |
277 | | - # ------- ._color.interpolation ---------------- |
278 | | - "ridgeplot._color.interpolation.ColorscaleInterpolants", |
279 | | - "ridgeplot._color.interpolation.SolidColormode", |
280 | | - # ------- ._kde -------------------------------- |
281 | | - "ridgeplot._kde.KDEPoints", |
282 | | - "ridgeplot._kde.KDEBandwidth", |
283 | | - # ------- ._missing ---------------------------- |
284 | | - "ridgeplot._missing.MISSING", |
285 | | - "ridgeplot._missing.MissingType", |
286 | | - # ------- ._types ------------------------------ |
287 | | - "ridgeplot._types.Color", |
288 | | - "ridgeplot._types.ColorScale", |
289 | | - "ridgeplot._types.NormalisationOption", |
290 | | - "ridgeplot._types.CollectionL1", |
291 | | - "ridgeplot._types.CollectionL2", |
292 | | - "ridgeplot._types.CollectionL3", |
293 | | - "ridgeplot._types.Float", |
294 | | - "ridgeplot._types.Int", |
295 | | - "ridgeplot._types.Numeric", |
296 | | - "ridgeplot._types.NumericT", |
297 | | - "ridgeplot._types.XYCoordinate", |
298 | | - "ridgeplot._types.DensityTrace", |
299 | | - "ridgeplot._types.DensitiesRow", |
300 | | - "ridgeplot._types.Densities", |
301 | | - "ridgeplot._types.ShallowDensities", |
302 | | - "ridgeplot._types.SamplesTrace", |
303 | | - "ridgeplot._types.SamplesRow", |
304 | | - "ridgeplot._types.Samples", |
305 | | - "ridgeplot._types.ShallowSamples", |
306 | | - "ridgeplot._types.TraceType", |
307 | | - "ridgeplot._types.TraceTypesArray", |
308 | | - "ridgeplot._types.ShallowTraceTypesArray", |
309 | | - "ridgeplot._types.LabelsArray", |
310 | | - "ridgeplot._types.ShallowLabelsArray", |
311 | | - "ridgeplot._types.SampleWeights", |
312 | | - "ridgeplot._types.SampleWeightsArray", |
313 | | - "ridgeplot._types.ShallowSampleWeightsArray", |
314 | | -} |
| 279 | +def _get_module_type_aliases(module_name: str) -> set[str]: |
| 280 | + module = import_module(module_name) |
| 281 | + source = inspect.getsource(module) |
| 282 | + tree = ast.parse(source) |
| 283 | + type_aliases = { |
| 284 | + f"{module_name}.{node.target.id}" |
| 285 | + for node in ast.walk(tree) |
| 286 | + # Handle annotated assignments: x: TypeAlias = ... |
| 287 | + if ( |
| 288 | + isinstance(node, ast.AnnAssign) |
| 289 | + and isinstance(node.target, ast.Name) |
| 290 | + and isinstance(node.annotation, ast.Name) |
| 291 | + and node.annotation.id == "TypeAlias" |
| 292 | + ) |
| 293 | + } |
| 294 | + return type_aliases |
| 295 | + |
| 296 | + |
| 297 | +def _get_package_type_aliases(package_name: str) -> set[str]: |
| 298 | + """Recursively extract TypeAlias definitions from a package.""" |
| 299 | + package = import_module(package_name) |
| 300 | + prefix = f"{package.__name__}." |
| 301 | + type_aliases = set() |
| 302 | + for _, modname, is_pkg in pkgutil.walk_packages(package.__path__, prefix): |
| 303 | + if is_pkg: |
| 304 | + type_aliases.update(_get_package_type_aliases(modname)) |
| 305 | + else: |
| 306 | + type_aliases.update(_get_module_type_aliases(modname)) |
| 307 | + return type_aliases |
| 308 | + |
| 309 | + |
| 310 | +_TYPE_ALIASES_FULLY_QUALIFIED = ( |
| 311 | + # Automatically extract all TypeAlias |
| 312 | + # definitions from the `ridgeplot` package |
| 313 | + _get_package_type_aliases(ridgeplot_pkg.__name__) |
| 314 | + | { |
| 315 | + # `MISSING` is a special case that is |
| 316 | + # widely referenced in type annotations |
| 317 | + "ridgeplot._missing.MISSING", |
| 318 | + # Same thing with the `NumericT` TypeVar |
| 319 | + "ridgeplot._types.NumericT", |
| 320 | + } |
| 321 | +) |
315 | 322 | for fq in _TYPE_ALIASES_FULLY_QUALIFIED: |
316 | 323 | module_name, _, type_name = fq.rpartition(".") |
317 | 324 | try: |
|
0 commit comments