Files
dbt-core/scripts/collect-dbt-contexts.py
Emily Rockman 8e800cee4c Consolidate pre-commit config (#12140)
* use single version of tools

* isort
2025-11-05 12:20:29 -05:00

114 lines
2.8 KiB
Python

#!/usr/bin/env python
import inspect
import json
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Union
from dbt.context.base import BaseContext
from dbt.context.providers import MacroContext, ModelContext
from dbt.context.target import TargetContext
from dbt_common.dataclass_schema import dbtClassMixin
CONTEXTS_MAP = {
"base": BaseContext,
"target": TargetContext,
"model": ModelContext,
"macro": MacroContext,
}
@dataclass
class ContextValue(dbtClassMixin):
name: str
value: str # a type description
doc: Optional[str]
@dataclass
class MethodArgument(dbtClassMixin):
name: str
value: str # a type description
@dataclass
class ContextMethod(dbtClassMixin):
name: str
args: List[MethodArgument]
result: str # a type description
doc: Optional[str]
@dataclass
class Unknown(dbtClassMixin):
name: str
value: str
doc: Optional[str]
ContextMember = Union[ContextValue, ContextMethod, Unknown]
def _get_args(func: inspect.Signature) -> Iterable[MethodArgument]:
found_first = False
for argname, arg in func.parameters.items():
if found_first is False and argname in {"self", "cls"}:
continue
if found_first is False:
found_first = True
yield MethodArgument(
name=argname,
value=inspect.formatannotation(arg.annotation),
)
def collect(cls):
values = []
for name, v in cls._context_members_.items():
attrname = cls._context_attrs_[name]
attrdef = getattr(cls, attrname)
doc = getattr(attrdef, "__doc__")
if inspect.isfunction(attrdef):
sig = inspect.signature(attrdef)
result = inspect.formatannotation(sig.return_annotation)
sig_good_part = ContextMethod(
name=name,
args=list(_get_args(sig)),
result=result,
doc=doc,
)
elif isinstance(attrdef, property):
sig = inspect.signature(attrdef.fget)
sig_txt = inspect.formatannotation(sig.return_annotation)
sig_good_part = ContextValue(name=name, value=sig_txt, doc=doc)
else:
sig_good_part = Unknown(name=name, value=repr(attrdef), doc=doc)
values.append(sig_good_part)
return values
@dataclass
class ContextCatalog(dbtClassMixin):
base: List[ContextMember]
target: List[ContextMember]
model: List[ContextMember]
macro: List[ContextMember]
schema: Dict[str, Any]
def main():
catalog = ContextCatalog(
base=collect(BaseContext),
target=collect(TargetContext),
model=collect(ModelContext),
macro=collect(MacroContext),
schema=ContextCatalog.json_schema(),
)
print(json.dumps(catalog.to_dict()))
if __name__ == "__main__":
main()