Repository: quora/pyanalyze
Pull request #585 (2023-01-16 17:57:51)

Add `class_attribute_transformers` hook
This is a more useful version of treat_class_attribute_as_any.

The use case I envision is to replace ORM fields with their values.



Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 1f0023f..0709721 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Add `class_attribute_transformers` hook (#585)
 - Support for PEP 702 (`@typing.deprecated`) (#578)
 - Simplify import handling; stop trying to import modules at type checking time (#566)
 - Suggest using keyword arguments on calls with too many positional arguments (#572)
diff --git a/pyanalyze/attributes.py b/pyanalyze/attributes.py
index 9dcffe2..efb4321 100644
--- a/pyanalyze/attributes.py
+++ b/pyanalyze/attributes.py
@@ -191,6 +191,36 @@ class TreatClassAttributeAsAny(PyObjectSequenceOption[_TCAA]):
         return any(func(val) for func in option_value)
 
 
+_CAT = Callable[[object], Optional[Tuple[Value, Value]]]
+
+
+class ClassAttributeTransformer(PyObjectSequenceOption[_CAT]):
+    """Transform certain class attributes.
+
+    Instances of this option are callables that take an object found among
+    a class's attributes and return either None (if the value should not
+    be transformed) or a pair of a get and set type. To disallow setting
+    the value, return :data:`pyanalyze.value.NO_RETURN_VALUE`.
+
+    If multiple transformers match an object, the first one is used.
+
+    TODO: The set type is currently ignored.
+
+    """
+
+    default_value: Sequence[_CAT] = []
+    name = "class_attribute_transformers"
+
+    @classmethod
+    def transform_attribute(cls, val: object, options: Options) -> Optional[Value]:
+        option_value = options.get_value_for(cls)
+        for transformer in option_value:
+            result = transformer(val)
+            if result is not None:
+                return result[0]
+        return None
+
+
 def _unwrap_value_from_subclass(result: Value, ctx: AttrContext) -> Value:
     if not isinstance(result, KnownValue) or ctx.skip_unwrap:
         return result
@@ -215,6 +245,11 @@ def _unwrap_value_from_subclass(result: Value, ctx: AttrContext) -> Value:
     elif TreatClassAttributeAsAny.should_treat_as_any(cls_val, ctx.options):
         return AnyValue(AnySource.error)
     else:
+        transformed = ClassAttributeTransformer.transform_attribute(
+            cls_val, ctx.options
+        )
+        if transformed is not None:
+            return transformed
         return KnownValue(cls_val)
 
 
@@ -328,6 +363,11 @@ def _unwrap_value_from_typed(result: Value, typ: type, ctx: AttrContext) -> Valu
     elif TreatClassAttributeAsAny.should_treat_as_any(cls_val, ctx.options):
         return AnyValue(AnySource.error)
     else:
+        transformed = ClassAttributeTransformer.transform_attribute(
+            cls_val, ctx.options
+        )
+        if transformed is not None:
+            return transformed
         return result
 
 


Test patch:
diff --git a/pyanalyze/test.toml b/pyanalyze/test.toml
index e0537a0..e551d77 100644
--- a/pyanalyze/test.toml
+++ b/pyanalyze/test.toml
@@ -10,3 +10,6 @@ functions_safe_to_call = [
     "pyanalyze.tests.make_simple_sequence",
     "pyanalyze.value.make_coro_type",
 ]
+class_attribute_transformers = [
+    "pyanalyze.test_config.transform_class_attribute"
+]
\ No newline at end of file
diff --git a/pyanalyze/test_attributes.py b/pyanalyze/test_attributes.py
index c44ff96..5a970a6 100644
--- a/pyanalyze/test_attributes.py
+++ b/pyanalyze/test_attributes.py
@@ -366,3 +366,15 @@ class TestHasAttrExtension(TestNameCheckVisitorBase):
             cls = X
             if hasattr(cls, "types"):  # E: value_always_true
                 assert_is_value(cls.types(), AnyValue(AnySource.unannotated))
+
+
+class TestClassAttributeTransformer(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test(self):
+        from pyanalyze.test_config import StringField
+
+        class Capybara:
+            foo = StringField()
+
+        def capybara(c: Capybara):
+            assert_is_value(c.foo, TypedValue(str))
diff --git a/pyanalyze/test_config.py b/pyanalyze/test_config.py
index c9387d8..6f705d1 100644
--- a/pyanalyze/test_config.py
+++ b/pyanalyze/test_config.py
@@ -4,7 +4,7 @@ Configuration file specific to tests.
 
 """
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, Optional, Tuple
 
 from . import tests, value
 from .arg_spec import ArgSpecCache
@@ -112,5 +112,19 @@ def unwrap_class(cls: type) -> type:
     return cls
 
 
+class StringField:
+    pass
+
+
+@used  # in test.toml
+def transform_class_attribute(
+    attr: object,
+) -> Optional[Tuple[value.Value, value.Value]]:
+    """Transforms a StringField attribute."""
+    if isinstance(attr, StringField):
+        return value.TypedValue(str), value.NO_RETURN_VALUE
+    return None
+
+
 CONFIG_PATH = Path(__file__).parent / "test.toml"
 TEST_OPTIONS = Options.from_option_list(config_file_path=CONFIG_PATH)