Repository: quora/pyanalyze
Pull request #669 (2023-08-23 23:53:49)

Add hook for overriding attributes on literals


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index ff048eb..4832bdd 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add hook for overriding the value inferred for attributes on
+  literals (#669)
 - Support classes that set a `__signature__` attribute to
   define their constructor signature (such as Pydantic models) (#665)
 - Declare support for Python 3.12. Not all new features in
diff --git a/pyanalyze/attributes.py b/pyanalyze/attributes.py
index 0e418f0..292ddb5 100644
--- a/pyanalyze/attributes.py
+++ b/pyanalyze/attributes.py
@@ -9,7 +9,7 @@ import sys
 import types
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Callable, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, ClassVar, Optional, Sequence, Tuple, Union
 
 import asynq
 import qcore
@@ -215,7 +215,7 @@ class ClassAttributeTransformer(PyObjectSequenceOption[_CAT]):
 
     """
 
-    default_value: Sequence[_CAT] = []
+    default_value: ClassVar[Sequence[_CAT]] = []
     name = "class_attribute_transformers"
 
     @classmethod
@@ -381,22 +381,49 @@ def _unwrap_value_from_typed(result: Value, typ: type, ctx: AttrContext) -> Valu
         return result
 
 
-def _get_attribute_from_known(obj: object, ctx: AttrContext) -> Value:
-    ctx.record_attr_read(type(obj))
+_KAH = Callable[[object, str], Optional[Value]]
 
-    if (obj is None or obj is NoneType) and ctx.should_ignore_none_attributes():
-        # This usually indicates some context is set to None
-        # in the module and initialized later.
-        return AnyValue(AnySource.error)
 
+def _default_transformer(obj: object, attr: str) -> Optional[Value]:
     # Type alias to Any
     if obj is Any:
         return AnyValue(AnySource.explicit)
 
     # Avoid generating huge Union type with the actual value
-    if obj is sys and ctx.attr == "modules":
+    if obj is sys and attr == "modules":
         return GenericValue(dict, [TypedValue(str), TypedValue(types.ModuleType)])
 
+    return None
+
+
+class KnownAttributeHook(PyObjectSequenceOption[_KAH]):
+    """Allows hooking into the inferred value for an attribute on a literal."""
+
+    default_value: ClassVar[Sequence[_KAH]] = [_default_transformer]
+    name = "known_attribute_hook"
+
+    @classmethod
+    def get_attribute(cls, obj: object, attr: str, options: Options) -> Optional[Value]:
+        option_value = options.get_value_for(cls)
+        for transformer in option_value:
+            result = transformer(obj, attr)
+            if result is not None:
+                return result
+        return None
+
+
+def _get_attribute_from_known(obj: object, ctx: AttrContext) -> Value:
+    ctx.record_attr_read(type(obj))
+
+    if (obj is None or obj is NoneType) and ctx.should_ignore_none_attributes():
+        # This usually indicates some context is set to None
+        # in the module and initialized later.
+        return AnyValue(AnySource.error)
+
+    hooked_value = KnownAttributeHook.get_attribute(obj, ctx.attr, ctx.options)
+    if hooked_value is not None:
+        return hooked_value
+
     result, _, _ = _get_attribute_from_mro(obj, ctx, on_class=True)
     if (
         isinstance(result, KnownValue)


Test patch:
diff --git a/pyanalyze/test.toml b/pyanalyze/test.toml
index d226448..22c456e 100644
--- a/pyanalyze/test.toml
+++ b/pyanalyze/test.toml
@@ -13,6 +13,9 @@ functions_safe_to_call = [
 class_attribute_transformers = [
     "pyanalyze.test_config.transform_class_attribute"
 ]
+known_attribute_hook = [
+    "pyanalyze.test_config.known_attribute_hook"
+]
 disallowed_imports = [
     "getopt",
     "email.quoprimime",
diff --git a/pyanalyze/test_attributes.py b/pyanalyze/test_attributes.py
index 2f36607..0d51dcf 100644
--- a/pyanalyze/test_attributes.py
+++ b/pyanalyze/test_attributes.py
@@ -70,6 +70,13 @@ class TestAttributes(TestNameCheckVisitorBase):
             assert_is_value(DefiniteCapybara().capybara_id, KnownValue(3))
             assert_is_value(DefiniteCapybara.capybara_id, KnownValue(3))
 
+    @assert_passes()
+    def test_known_value_hook(self):
+        from pyanalyze.test_config import SPECIAL_STRING
+
+        def capybara():
+            assert_is_value(SPECIAL_STRING.special, KnownValue("special"))
+
     @assert_passes()
     def test_generic(self):
         from typing import Generic, TypeVar
diff --git a/pyanalyze/test_config.py b/pyanalyze/test_config.py
index 6f705d1..02f32c7 100644
--- a/pyanalyze/test_config.py
+++ b/pyanalyze/test_config.py
@@ -126,5 +126,15 @@ def transform_class_attribute(
     return None
 
 
+SPECIAL_STRING = "capybara!"
+
+
+@used  # in test.toml
+def known_attribute_hook(obj: object, attr: str) -> Optional[value.Value]:
+    if obj == SPECIAL_STRING and attr == "special":
+        return value.KnownValue("special")
+    return None
+
+
 CONFIG_PATH = Path(__file__).parent / "test.toml"
 TEST_OPTIONS = Options.from_option_list(config_file_path=CONFIG_PATH)