Implement _StrPromise attribute hook.

This implements an attribute hook that provides type information for
methods that are available on `builtins.str` for `_StrPromise` except
the supported operators. This allows us to avoid copying stubs from the
builtins for all supported methods on `str`.

Signed-off-by: Zixuan James Li <[email protected]>

Author: PIG208

django-stubs/utils/functional.pyi

@@ -40,6 +40,8 @@ class _StrPromise(Promise, Sequence[str]):
     def __mod__(self, __x: Any) -> str: ...
     def __mul__(self, __n: SupportsIndex) -> str: ...
     def __rmul__(self, __n: SupportsIndex) -> str: ...
+    # Mypy requires this for the attribute hook to take effect
+    def __getattribute__(self, __name: str) -> Any: ...
 
 _C = TypeVar("_C", bound=Callable)
 

mypy_django_plugin/lib/fullnames.py

@@ -41,3 +41,5 @@
 F_EXPRESSION_FULLNAME = "django.db.models.expressions.F"
 
 ANY_ATTR_ALLOWED_CLASS_FULLNAME = "django_stubs_ext.AnyAttrAllowed"
+
+STR_PROMISE_FULLNAME = "django.utils.functional._StrPromise"

mypy_django_plugin/main.py

@@ -22,6 +22,7 @@
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.transformers import fields, forms, init_create, meta, querysets, request, settings
+from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute
 from mypy_django_plugin.transformers.managers import (
     create_new_manager_class_from_from_queryset_method,
     resolve_manager_method,
@@ -285,6 +286,9 @@ def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeConte
         ):
             return resolve_manager_method
 
+        if info and info.has_base(fullnames.STR_PROMISE_FULLNAME):
+            return resolve_str_promise_attribute
+
         return None
 
     def get_type_analyze_hook(self, fullname: str) -> Optional[Callable[[AnalyzeTypeContext], MypyType]]:

mypy_django_plugin/transformers/functional.py

@@ -0,0 +1,40 @@
+from mypy.errorcodes import ATTR_DEFINED
+from mypy.nodes import CallExpr, MemberExpr
+from mypy.plugin import AttributeContext
+from mypy.types import AnyType, CallableType
+from mypy.types import Type as MypyType
+from mypy.types import TypeOfAny
+
+from mypy_django_plugin.lib import helpers
+
+
+def resolve_str_promise_attribute(ctx: AttributeContext) -> MypyType:
+    if isinstance(ctx.context, MemberExpr):
+        method_name = ctx.context.name
+    elif isinstance(ctx.context, CallExpr) and isinstance(ctx.context.callee, MemberExpr):
+        method_name = ctx.context.callee.name
+    else:
+        ctx.api.fail(f'Cannot resolve the attribute of "{ctx.type}"', ctx.context, code=ATTR_DEFINED)
+        return AnyType(TypeOfAny.from_error)
+
+    str_info = helpers.lookup_fully_qualified_typeinfo(helpers.get_typechecker_api(ctx), f"builtins.str")
+    assert str_info is not None
+    method = str_info.get(method_name)
+
+    if method is None or method.type is None:
+        ctx.api.fail(f'"{ctx.type}" has no attribute "{method_name}"', ctx.context, code=ATTR_DEFINED)
+        return AnyType(TypeOfAny.from_error)
+
+    if isinstance(method.type, CallableType):
+        # The proxied str methods are only meant to be used as instance methods.
+        # We need to drop the first `self` argument in them.
+        assert method.type.arg_names[0] == "self"
+        return method.type.copy_modified(
+            arg_kinds=method.type.arg_kinds[1:],
+            arg_names=method.type.arg_names[1:],
+            arg_types=method.type.arg_types[1:],
+        )
+    else:
+        # Not possible with `builtins.str`, but we have error handling for this anyway.
+        ctx.api.fail(f'"{method_name}" on "{ctx.type}" is not a method', ctx.context)
+        return AnyType(TypeOfAny.from_error)

tests/typecheck/utils/test_functional.yml

@@ -16,3 +16,34 @@
       f = Foo()
       reveal_type(f.attr)        # N: Revealed type is "builtins.list[builtins.str]"
       f.attr.name                # E: "List[str]" has no attribute "name"
+
+-   case: str_promise_proxy
+    main:  |
+      from typing import Union
+
+      from django.utils.functional import Promise, lazystr, _StrPromise
+
+      s = lazystr("asd")
+
+      reveal_type(s)                 # N: Revealed type is "django.utils.functional._StrPromise"
+
+      reveal_type(s.format("asd"))   # N: Revealed type is "builtins.str"
+      reveal_type(s.capitalize())    # N: Revealed type is "builtins.str"
+      reveal_type(s.swapcase)        # N: Revealed type is "def () -> builtins.str"
+      reveal_type(s.__getnewargs__)  # N: Revealed type is "def () -> Tuple[builtins.str]"
+      s.nonsense                     # E: "django.utils.functional._StrPromise" has no attribute "nonsense"
+      f: Union[_StrPromise, str]
+      reveal_type(f.format("asd"))   # N: Revealed type is "builtins.str"
+
+      reveal_type(s + "bar")         # N: Revealed type is "builtins.str"
+      reveal_type("foo" + s)         # N: Revealed type is "Any"
+      reveal_type(s % "asd")         # N: Revealed type is "builtins.str"
+
+      def foo(content: str) -> None:
+        ...
+
+      def bar(content: Promise) -> None:
+        ...
+
+      foo(s)                         # E: Argument 1 to "foo" has incompatible type "_StrPromise"; expected "str"
+      bar(s)