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,20 @@
-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 .base import Storage
 
+@type_check_only
+class _StorageConfig(TypedDict):
+    BACKEND: str
+    OPTIONS: dict[str, Any] | None
+
 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,21 @@ 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":
+            return [
+                self._new_dependency(self.django_context.settings.STORAGES["staticfiles"]["BACKEND"].rsplit(".", 1)[0])
+            ]
+
+        # for settings.STORAGES
+        elif file.fullname == "django.core.files.storage":
+            return [
+                self._new_dependency(s["BACKEND"].rsplit(".", 1)[0])
+                for s in self.django_context.settings.STORAGES.values()
+            ]
+
         # 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 +214,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 +315,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:
@@ -312,8 +333,11 @@ 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
+        # 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"] = [s["BACKEND"] for s in self.django_context.settings.STORAGES.values()]
+        # 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,47 @@
+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, get_proper_type
+from mypy.types import Type as MypyType
+
+from mypy_django_plugin.django.context import DjangoContext
+from mypy_django_plugin.lib import helpers
+
+
+def get_storage(ctx: AnalyzeTypeContext, alias: str, django_context: DjangoContext) -> MypyType:
+    assert isinstance(ctx.api, TypeAnalyser)
+    assert isinstance(ctx.api.api, SemanticAnalyzer)
+
+    fullname = django_context.settings.STORAGES[alias]["BACKEND"]
+    type_info = helpers.lookup_fully_qualified_typeinfo(ctx.api.api, fullname)
+
+    if type_info is None:
+        if not ctx.api.api.final_iteration:
+            ctx.api.api.defer()
+            return PlaceholderType(fullname=fullname, args=[], line=ctx.context.line)
+        else:
+            ctx.api.fail(f'Cannot resolve storage backend "{fullname}"', ctx.context)
+            return ctx.type
+
+    assert type_info
+    return Instance(type_info, [])
+
+
+def extract_proper_type_for_getitem(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
+    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_literal.value, str)
+        ):
+            fullname = django_context.settings.STORAGES[alias_literal.value]["BACKEND"]
+            type_info = helpers.lookup_fully_qualified_typeinfo(ctx.api, fullname)
+            assert type_info
+            return Instance(type_info, [])
+
+    return ctx.default_return_type

tests/typecheck/contrib/staticfiles/test_storage.yml

@@ -0,0 +1,5 @@
+- case: test_staticfiles_storage_defaults
+  main: |
+    from django.contrib.staticfiles.storage import staticfiles_storage
+
+    reveal_type(staticfiles_storage)  # N: Revealed type is "django.contrib.staticfiles.storage.StaticFilesStorage"

tests/typecheck/core/test_storage.yml

@@ -0,0 +1,39 @@
+- case: test_storage_defaults
+  main: |
+    from django.core.files.storage import default_storage, storages
+
+    reveal_type(default_storage)  # N: Revealed type is "django.core.files.storage.filesystem.FileSystemStorage"
+    reveal_type(storages["default"])  # N: Revealed type is "django.core.files.storage.filesystem.FileSystemStorage"
+    reveal_type(storages["staticfiles"])  # N: Revealed type is "django.contrib.staticfiles.storage.StaticFilesStorage"
+
+- case: test_custom_storages
+  main: |
+    from django.core.files.storage import default_storage, storages
+
+    reveal_type(default_storage)  # N: Revealed type is "myapp.storage.MyDefaultStorage"
+    reveal_type(storages["default"])  # N: Revealed type is "myapp.storage.MyDefaultStorage"
+    reveal_type(storages["custom"])  # N: Revealed type is "myapp.storage.MyStorage"
+    reveal_type(storages["staticfiles"])  # N: Revealed type is "django.contrib.staticfiles.storage.StaticFilesStorage"
+
+  custom_settings: |
+    from django.conf.global_settings import STORAGES as DEFAULT_STORAGES
+
+    STORAGES = {
+        **DEFAULT_STORAGES,
+        "default": {"BACKEND": "myapp.storage.MyDefaultStorage"},
+        "custom": {
+            "BACKEND": "myapp.storage.MyStorage",
+            "OPTIONS": {"option_enabled": False, "key": "test"},
+        }
+    }
+
+  files:
+    - path: myapp/storage.py
+      content: |
+        from django.core.files.storage import Storage
+
+        class MyDefaultStorage(Storage):
+            pass
+
+        class MyStorage(Storage):
+            pass