From 085e30bb5656741464e753ba0b92e6942d434c82 Mon Sep 17 00:00:00 2001 From: Quigley Malcolm Date: Wed, 30 Aug 2023 16:00:27 -0700 Subject: [PATCH] Improve typing `BaseContext` and functions In addition to just adding parameter typing and return typing to `BaseContext` functions. We also declared `_context_members_` and `_context_attrs_` as properites of `BaseContext` this was necessary because they're being accessed in the classes functions. However, because they were being indirectly instantiated by the metaclass `ContextMeta`, the properties weren't actually known to exist. By adding declaring the properties on the `BaseContext`, we let mypy know they exist. --- core/dbt/context/base.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/core/dbt/context/base.py b/core/dbt/context/base.py index a5717c021ec..6eacfdf9b78 100644 --- a/core/dbt/context/base.py +++ b/core/dbt/context/base.py @@ -181,13 +181,17 @@ def __call__(self, var_name: str, default: Any = _VAR_NOTSET) -> Any: class BaseContext(metaclass=ContextMeta): + # Set by ContextMeta + _context_members_: Dict[str, Any] + _context_attrs_: Dict[str, Any] + # subclass is TargetContext - def __init__(self, cli_vars): - self._ctx = {} - self.cli_vars = cli_vars - self.env_vars = {} + def __init__(self, cli_vars: Dict[str, Any]): + self._ctx: Dict[str, Any] = {} + self.cli_vars: Dict[str, Any] = cli_vars + self.env_vars: Dict[str, Any] = {} - def generate_builtins(self): + def generate_builtins(self) -> Dict[str, Any]: builtins: Dict[str, Any] = {} for key, value in self._context_members_.items(): if hasattr(value, "__get__"): @@ -197,7 +201,7 @@ def generate_builtins(self): return builtins # no dbtClassMixin so this is not an actual override - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: self._ctx["context"] = self._ctx builtins = self.generate_builtins() self._ctx["builtins"] = builtins