|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import functools |
| 4 | +import os.path |
4 | 5 | from collections.abc import Callable |
5 | 6 | from typing import Any |
6 | 7 |
|
7 | 8 | from mypy.build import PRI_MYPY |
8 | 9 | from mypy.errorcodes import ATTR_DEFINED |
| 10 | +from mypy.lookup import lookup_fully_qualified |
9 | 11 | from mypy.messages import format_type |
10 | 12 | from mypy.nodes import ARG_POS, DictExpr, ListExpr, MypyFile, StrExpr, TypeInfo |
11 | 13 | from mypy.plugin import ( |
|
16 | 18 | MethodContext, |
17 | 19 | MethodSigContext, |
18 | 20 | Plugin, |
| 21 | + ReportConfigContext, |
19 | 22 | SemanticAnalyzerPluginInterface, |
20 | 23 | ) |
21 | 24 | from mypy.plugins.common import add_attribute_to_class, add_method_to_class |
@@ -160,6 +163,40 @@ def _adjust_request_members(ctx: ClassDefContext) -> None: |
160 | 163 | add_attribute_to_class(ctx.api, ctx.cls, "auth", _request_auth_tp(ctx.api)) |
161 | 164 |
|
162 | 165 |
|
| 166 | +_REQUEST_USER_FULLNAME = "rest_framework.request.Request.user" |
| 167 | +_REQUEST_USER_TPS = ( |
| 168 | + "sentry.users.models.user.User", |
| 169 | + "django.contrib.auth.models.AnonymousUser", |
| 170 | + "sentry.service_accounts.principal.ServiceAccountPrincipal", |
| 171 | +) |
| 172 | +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| 173 | +_REQUEST_USER_ALLOWLIST = os.path.join(os.path.dirname(__file__), "request_user_allowlist.txt") |
| 174 | + |
| 175 | + |
| 176 | +@functools.cache |
| 177 | +def _request_user_allowlist() -> frozenset[str]: |
| 178 | + with open(_REQUEST_USER_ALLOWLIST) as f: |
| 179 | + return frozenset(line.strip() for line in f if line.strip()) |
| 180 | + |
| 181 | + |
| 182 | +def _is_request_user_allowlisted(path: str) -> bool: |
| 183 | + return os.path.relpath(os.path.abspath(path), _REPO_ROOT) in _request_user_allowlist() |
| 184 | + |
| 185 | + |
| 186 | +def _request_user_attribute(ctx: AttributeContext) -> Type: |
| 187 | + # sentry.auth.principal.set_request_principal can set a ServiceAccountPrincipal |
| 188 | + if _is_request_user_allowlisted(ctx.api.path): |
| 189 | + return ctx.default_attr_type |
| 190 | + # named_generic_type walks parent packages, which don't always list these submodules |
| 191 | + modules = ctx.api.modules # type: ignore[attr-defined] |
| 192 | + items: list[Type] = [] |
| 193 | + for tp in _REQUEST_USER_TPS: |
| 194 | + st = lookup_fully_qualified(tp, modules, raise_on_missing=True) |
| 195 | + assert st is not None and isinstance(st.node, TypeInfo), tp |
| 196 | + items.append(Instance(st.node, ())) |
| 197 | + return UnionType(items) |
| 198 | + |
| 199 | + |
163 | 200 | def _add_principal_members_to_anonymous_user(ctx: ClassDefContext) -> None: |
164 | 201 | if ctx.cls.name == "AnonymousUser": |
165 | 202 | # added by sentry.auth.principal.patch_request_principal |
@@ -550,12 +587,24 @@ def get_attribute_hook(self, fullname: str) -> Callable[[AttributeContext], Type |
550 | 587 | if fullname.startswith("sentry.utils.lazy_service_wrapper.LazyServiceWrapper."): |
551 | 588 | _, attr = fullname.rsplit(".", 1) |
552 | 589 | return functools.partial(_lazy_service_wrapper_attribute, attr=attr) |
| 590 | + elif fullname == _REQUEST_USER_FULLNAME: |
| 591 | + return _request_user_attribute |
553 | 592 | else: |
554 | 593 | return None |
555 | 594 |
|
| 595 | + def report_config_data(self, ctx: ReportConfigContext) -> bool: |
| 596 | + # recheck a cached module when it joins or leaves the request.user allowlist |
| 597 | + return _is_request_user_allowlisted(ctx.path) |
| 598 | + |
556 | 599 | def get_additional_deps(self, file: MypyFile) -> list[tuple[int, str, int]]: |
557 | | - if file.fullname in {"django.http", "django.http.request", "rest_framework.request"}: |
| 600 | + if file.fullname in {"django.http", "django.http.request"}: |
558 | 601 | return [(PRI_MYPY, "sentry.auth.services.auth.model", -1)] |
| 602 | + elif file.fullname == "rest_framework.request": |
| 603 | + return [ |
| 604 | + (PRI_MYPY, "sentry.auth.services.auth.model", -1), |
| 605 | + (PRI_MYPY, "sentry.users.models.user", -1), |
| 606 | + (PRI_MYPY, "sentry.service_accounts.principal", -1), |
| 607 | + ] |
559 | 608 | else: |
560 | 609 | return [] |
561 | 610 |
|
|
0 commit comments