add mypy plugin support for Django's storage framework

This change correctly resolves
``django.core.files.storage.default_storage`` and
``django.contrib.staticfiles.storage.staticfiles_storage`` which are
configured in the project's Django settings.

When ``django.core.files.storage.storages``, which has a dict like
interface, is accessed with a literal value, then the type of the
appropriate storage is returned.

Author: terencehonles

django-stubs/contrib/staticfiles/finders.pyi

@@ -2,7 +2,7 @@ from collections.abc import Iterable, Iterator, Sequence
 from typing import Any, Literal, overload
 
 from django.core.checks.messages import CheckMessage
-from django.core.files.storage import FileSystemStorage, Storage
+from django.core.files.storage import FileSystemStorage, Storage, _DefaultStorage
 
 searched_locations: Any
 
@@ -16,7 +16,7 @@ class BaseFinder:
 
 class FileSystemFinder(BaseFinder):
     locations: list[tuple[str, str]]
-    storages: dict[str, Any]
+    storages: dict[str, FileSystemStorage]
     def __init__(self, app_names: Sequence[str] | None = None, *args: Any, **kwargs: Any) -> None: ...
     def find_location(self, root: str, path: str, prefix: str | None = None) -> str | None: ...
     @overload
@@ -48,7 +48,7 @@ class BaseStorageFinder(BaseFinder):
     def list(self, ignore_patterns: Iterable[str] | None) -> Iterable[Any]: ...
 
 class DefaultStorageFinder(BaseStorageFinder):
-    storage: Storage
+    storage: _DefaultStorage
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
 
 @overload

django-stubs/contrib/staticfiles/management/commands/collectstatic.pyi

@@ -1,5 +1,6 @@
 from typing import Any
 
+from django.contrib.staticfiles.storage import _ConfiguredStorage
 from django.core.files.storage import Storage
 from django.core.management.base import BaseCommand
 from django.utils.functional import cached_property
@@ -9,7 +10,7 @@ class Command(BaseCommand):
     symlinked_files: Any
     unmodified_files: Any
     post_processed_files: Any
-    storage: Any
+    storage: _ConfiguredStorage
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
     @cached_property
     def local(self) -> bool: ...

django-stubs/contrib/staticfiles/storage.pyi

@@ -52,4 +52,8 @@ class ManifestFilesMixin(HashedFilesMixin):
 class ManifestStaticFilesStorage(ManifestFilesMixin, StaticFilesStorage): ...  # type: ignore[misc]
 class ConfiguredStorage(LazyObject): ...
 
-staticfiles_storage: Storage
+# This is our "placeholder" type the mypy plugin refines to configured
+# 'STORAGES["staticfiles"]["BACKEND"]' wherever it is used as a type.
+_ConfiguredStorage: TypeAlias = ConfiguredStorage
+
+staticfiles_storage: _ConfiguredStorage

django-stubs/core/files/storage/__init__.pyi

@@ -1,3 +1,5 @@
+from typing import TypeAlias
+
 from django.utils.functional import LazyObject
 
 from .base import Storage
@@ -18,6 +20,9 @@ __all__ = (
 
 class DefaultStorage(LazyObject): ...
 
+# This is our "placeholder" type the mypy plugin refines to configured
+# 'STORAGES["default"]["BACKEND"]' wherever it is used as a type.
+_DefaultStorage: TypeAlias = DefaultStorage
+
 storages: StorageHandler
-# default_storage is actually an instance of DefaultStorage, but it proxies through to a Storage
-default_storage: Storage
+default_storage: _DefaultStorage

django-stubs/core/files/storage/handler.pyi

@@ -1,15 +1,21 @@
-from typing import Any
+from typing import Any, TypedDict, type_check_only
 
 from django.core.exceptions import ImproperlyConfigured
 from django.utils.functional import cached_property
+from typing_extensions import NotRequired
 
 from .base import Storage
 
+@type_check_only
+class _StorageConfig(TypedDict):
+    BACKEND: str
+    OPTIONS: NotRequired[dict[str, Any]]
+
 class InvalidStorageError(ImproperlyConfigured): ...
 
 class StorageHandler:
-    def __init__(self, backends: dict[str, Storage] | None = None) -> None: ...
+    def __init__(self, backends: dict[str, _StorageConfig] | None = None) -> None: ...
     @cached_property
-    def backends(self) -> dict[str, Storage]: ...
+    def backends(self) -> dict[str, _StorageConfig]: ...
     def __getitem__(self, alias: str) -> Storage: ...
-    def create_storage(self, params: dict[str, Any]) -> Storage: ...
+    def create_storage(self, params: _StorageConfig) -> Storage: ...

mypy_django_plugin/lib/fullnames.py

@@ -11,6 +11,7 @@
 MANYTOMANY_FIELD_FULLNAME = "django.db.models.fields.related.ManyToManyField"
 DUMMY_SETTINGS_BASE_CLASS = "django.conf._DjangoConfLazyObject"
 AUTH_USER_MODEL_FULLNAME = "django.conf.settings.AUTH_USER_MODEL"
+STORAGE_HANDLER_CLASS_FULLNAME = "django.core.files.storage.handler.StorageHandler"
 
 QUERYSET_CLASS_FULLNAME = "django.db.models.query.QuerySet"
 BASE_MANAGER_CLASS_FULLNAME = "django.db.models.manager.BaseManager"

mypy_django_plugin/main.py

@@ -35,6 +35,7 @@
     orm_lookups,
     querysets,
     settings,
+    storage,
 )
 from mypy_django_plugin.transformers.auth import get_user_model
 from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute
@@ -98,8 +99,29 @@ def get_additional_deps(self, file: MypyFile) -> list[tuple[int, str, int]]:
         if file.fullname == "django.conf" and self.django_context.django_settings_module:
             return [self._new_dependency(self.django_context.django_settings_module, PRI_MED)]
 
+        # for settings.STORAGES["staticfiles"]
+        if (
+            file.fullname == "django.contrib.staticfiles.storage"
+            and isinstance(storage_config := self.django_context.settings.STORAGES.get("staticfiles"), dict)
+            and isinstance(storage_backend := storage_config.get("BACKEND"), str)
+            and "." in storage_backend
+        ):
+            return [self._new_dependency(storage_backend.rsplit(".", 1)[0])]
+
+        # for settings.STORAGES
+        elif file.fullname == "django.core.files.storage":
+            return [
+                self._new_dependency(storage_backend.rsplit(".", 1)[0])
+                for storage_config in self.django_context.settings.STORAGES.values()
+                if (
+                    isinstance(storage_config, dict)
+                    and isinstance(storage_backend := storage_config.get("BACKEND"), str)
+                    and "." in storage_backend
+                )
+            ]
+
         # for values / values_list
-        if file.fullname == "django.db.models":
+        elif file.fullname == "django.db.models":
             return [self._new_dependency("typing"), self._new_dependency("django_stubs_ext")]
 
         # for `get_user_model()`
@@ -200,6 +222,9 @@ def get_method_hook(self, fullname: str) -> Callable[[MethodContext], MypyType]
             }
             return hooks.get(class_fullname)
 
+        elif method_name == "__getitem__" and class_fullname == fullnames.STORAGE_HANDLER_CLASS_FULLNAME:
+            return partial(storage.extract_proper_type_for_getitem, django_context=self.django_context)
+
         if method_name in self.manager_and_queryset_method_hooks:
             info = self._get_typeinfo_or_none(class_fullname)
             if info and helpers.has_any_of_bases(
@@ -298,6 +323,10 @@ def get_type_analyze_hook(self, fullname: str) -> Callable[[AnalyzeTypeContext],
             return partial(handle_annotated_type, fullname=fullname)
         elif fullname == "django.contrib.auth.models._User":
             return partial(get_user_model, django_context=self.django_context)
+        elif fullname == "django.contrib.staticfiles.storage._ConfiguredStorage":
+            return partial(storage.get_storage, alias="staticfiles", django_context=self.django_context)
+        elif fullname == "django.core.files.storage._DefaultStorage":
+            return partial(storage.get_storage, alias="default", django_context=self.django_context)
         return None
 
     def get_dynamic_class_hook(self, fullname: str) -> Callable[[DynamicClassDefContext], None] | None:
@@ -311,9 +340,20 @@ def get_dynamic_class_hook(self, fullname: str) -> Callable[[DynamicClassDefCont
 
     def report_config_data(self, ctx: ReportConfigContext) -> dict[str, Any]:
         # Cache would be cleared if any settings do change.
-        extra_data = {}
-        # In all places we use '_User' alias as a type we want to clear cache if
-        # AUTH_USER_MODEL setting changes
+        extra_data: dict[str, Any] = {}
+        # In all places we use '_DefaultStorage' or '_ConfiguredStorage' aliases as a type we want to clear the cache
+        # if STORAGES setting changes
+        if ctx.id.startswith("django.contrib.staticfiles") or ctx.id.startswith("django.core.files.storage"):
+            extra_data["STORAGES"] = [
+                storage_backend
+                for storage_config in self.django_context.settings.STORAGES.values()
+                if (
+                    isinstance(storage_config, dict)
+                    and isinstance(storage_backend := storage_config.get("BACKEND"), str)
+                    and "." in storage_backend
+                )
+            ]
+        # In all places we use '_User' alias as a type we want to clear the cache if AUTH_USER_MODEL setting changes
         if ctx.id.startswith("django.contrib.auth") or ctx.id in {"django.http.request", "django.test.client"}:
             extra_data["AUTH_USER_MODEL"] = self.django_context.settings.AUTH_USER_MODEL
         return self.plugin_config.to_json(extra_data)

mypy_django_plugin/transformers/storage.py

@@ -0,0 +1,74 @@
+from mypy.checker import TypeChecker
+from mypy.plugin import AnalyzeTypeContext, MethodContext
+from mypy.semanal import SemanticAnalyzer
+from mypy.typeanal import TypeAnalyser
+from mypy.types import Instance, PlaceholderType, UninhabitedType, get_proper_type
+from mypy.types import Type as MypyType
+from mypy.typevars import fill_typevars
+
+from mypy_django_plugin.django.context import DjangoContext
+from mypy_django_plugin.lib import helpers
+
+
+def get_storage_backend(alias: str, django_context: DjangoContext) -> str | None:
+    "Defensively look for a settings.STORAGES by its alias."
+
+    try:
+        fullname = django_context.settings.STORAGES[alias]["BACKEND"]
+        if not isinstance(fullname, str) or "." not in fullname:
+            return None
+
+        return fullname
+    except (KeyError, TypeError):
+        return None
+
+
+def get_storage(ctx: AnalyzeTypeContext, alias: str, django_context: DjangoContext) -> MypyType:
+    """
+    Get a storage type by its alias, but do not fail if it cannot be found since this is resolving an internal type-var,
+    and errors would be reported in the type stubs.
+    """
+
+    assert isinstance(ctx.api, TypeAnalyser)
+    assert isinstance(ctx.api.api, SemanticAnalyzer)
+
+    if fullname := get_storage_backend(alias, django_context):
+        if type_info := helpers.lookup_fully_qualified_typeinfo(ctx.api.api, fullname):
+            return fill_typevars(type_info)
+
+        if not ctx.api.api.final_iteration:
+            ctx.api.api.defer()
+            return PlaceholderType(fullname=fullname, args=[], line=ctx.context.line)
+
+    return ctx.type
+
+
+def extract_proper_type_for_getitem(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
+    """
+    Provide type information for `StorageHandler.__getitem__` when providing a literal value.
+    """
+
+    assert isinstance(ctx.api, TypeChecker)
+
+    if ctx.arg_types:
+        alias_type = get_proper_type(ctx.arg_types[0][0])
+
+        if (
+            isinstance(alias_type, Instance)
+            and (alias_literal := alias_type.last_known_value)
+            and isinstance(alias := alias_literal.value, str)
+        ):
+            if alias not in django_context.settings.STORAGES:
+                ctx.api.fail(f'Could not find config for "{alias}" in settings.STORAGES.', ctx.context)
+
+            elif fullname := get_storage_backend(alias, django_context):
+                type_info = helpers.lookup_fully_qualified_typeinfo(ctx.api, fullname)
+                assert type_info
+                return fill_typevars(type_info)
+
+            else:
+                ctx.api.fail(f'"{alias}" in settings.STORAGES is improperly configured.', ctx.context)
+
+            return UninhabitedType()
+
+    return ctx.default_return_type

tests/assert_type/contrib/staticfiles/test_storage.py

@@ -0,0 +1,8 @@
+from django.contrib.staticfiles.storage import ConfiguredStorage, StaticFilesStorage, staticfiles_storage
+from typing_extensions import assert_type
+
+# The plugin can figure out what these are (but pyright can't):
+assert_type(staticfiles_storage, StaticFilesStorage)  # pyright: ignore[reportAssertTypeFailure]
+
+# what pyright thinks these are:
+assert_type(staticfiles_storage, ConfiguredStorage)  # mypy: ignore[assert-type]

tests/assert_type/core/files/test_storage.py

@@ -0,0 +1,10 @@
+from django.core.files.storage import DefaultStorage, FileSystemStorage, Storage, default_storage, storages
+from typing_extensions import assert_type
+
+# The plugin can figure out what these are (but pyright can't):
+assert_type(default_storage, FileSystemStorage)  # pyright: ignore[reportAssertTypeFailure]
+assert_type(storages["default"], FileSystemStorage)  # pyright: ignore[reportAssertTypeFailure]
+
+# what pyright thinks these are:
+assert_type(default_storage, DefaultStorage)  # mypy: ignore[assert-type]
+assert_type(storages["default"], Storage)  # mypy: ignore[assert-type]