Preserve Annotated annotations on access to methods of literals (#541)

Author: JelleZijlstra

docs/changelog.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Preserve `Annotated` annotations on access to methods of
+  literals (#541)
 - `allow_call` callables are now also called if the arguments
   are literals wrapped in `Annotated` (#540)
 - Support Python 3.11 (#537)

pyanalyze/attributes.py

@@ -348,6 +348,16 @@ def _get_attribute_from_known(obj: object, ctx: AttrContext) -> Value:
         return GenericValue(dict, [TypedValue(str), TypedValue(types.ModuleType)])
 
     result, _, _ = _get_attribute_from_mro(obj, ctx, on_class=True)
+    if (
+        isinstance(result, KnownValue)
+        and (
+            safe_isinstance(result.val, types.MethodType)
+            or safe_isinstance(result.val, types.BuiltinFunctionType)
+            and result.val.__self__ is obj
+        )
+        and isinstance(ctx.root_value, AnnotatedValue)
+    ):
+        result = UnboundMethodValue(ctx.attr, ctx.root_composite)
     if safe_isinstance(obj, type):
         result = set_self(result, TypedValue(obj))
     if isinstance(obj, (types.ModuleType, type)):

pyanalyze/test_attributes.py

@@ -4,8 +4,10 @@
 from .test_name_check_visitor import TestNameCheckVisitorBase
 from .test_node_visitor import assert_passes, only_before
 from .value import (
+    AnnotatedValue,
     AnySource,
     AnyValue,
+    UnboundMethodValue,
     assert_is_value,
     GenericValue,
     KnownValue,
@@ -170,6 +172,31 @@ def test(x: Union[Capybara, Paca]) -> None:
                 x.attr, MultiValuedValue([TypedValue(int), TypedValue(str)])
             )
 
+    @assert_passes()
+    def test_annotated_known(self):
+        from typing_extensions import Annotated, Literal
+        from pyanalyze.extensions import LiteralOnly
+        from pyanalyze.stacked_scopes import Composite, VarnameWithOrigin
+        from pyanalyze.value import CustomCheckExtension
+        from qcore.testing import Anything
+
+        origin = VarnameWithOrigin("encoding", Anything)  # E: incompatible_argument
+
+        def capybara():
+            encoding: Annotated[Literal["ascii"], LiteralOnly()] = "ascii"
+            assert_is_value(
+                encoding.encode,
+                UnboundMethodValue(
+                    "encode",
+                    Composite(
+                        AnnotatedValue(
+                            KnownValue("ascii"), [CustomCheckExtension(LiteralOnly())]
+                        ),
+                        origin,
+                    ),
+                ),
+            )
+
     @assert_passes()
     def test_optional_operation(self):
         from typing import Optional