WIP: Add module_only option to install_cache()

1 files changed, 67 insertions, 5 deletions

diff --git a/requests_cache/patcher.py b/requests_cache/patcher.py
index 406ee45..f32ef9e 100644
--- a/requests_cache/patcher.py
+++ b/requests_cache/patcher.py
@@ -7,9 +7,10 @@
    :functions-only:
    :nosignatures:
 """
+import inspect
 from contextlib import contextmanager
 from logging import getLogger
-from typing import Optional, Type
+from typing import TYPE_CHECKING, List, Optional, Type
 from warnings import warn
 
 import requests
@@ -19,10 +20,54 @@ from .session import CachedSession, OriginalSession
 
 logger = getLogger(__name__)
 
+if TYPE_CHECKING:
+    MIXIN_BASE = CachedSession
+else:
+    MIXIN_BASE = object
+
+
+class ModuleCacheMixin(MIXIN_BASE):
+    """Session mixin that optionally caches requests only if sent from specific modules"""
+
+    def __init__(
+        self, *args, module_only: bool = False, modules: Optional[List[str]] = None, **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.modules = modules or []
+        self.module_only = module_only
+
+    def request(self, *args, **kwargs):
+        if self._is_module_enabled():
+            return super().request(*args, **kwargs)
+        else:
+            return OriginalSession.request(self, *args, **kwargs)
+
+    def _is_module_enabled(self) -> bool:
+        if not self.module_only:
+            return True
+        return _calling_module(back=3) in self.modules
+
+    def enable_module(self):
+        self.modules.append(_calling_module())
+
+    def disable_module(self):
+        try:
+            self.modules.remove(_calling_module())
+        except ValueError:
+            pass
+
+
+def _calling_module(back: int = 2) -> str:
+    """Get the name of the module ``back`` frames up in the call stack"""
+    frame = inspect.stack()[back].frame
+    module = inspect.getmodule(frame)
+    return getattr(module, '__name__', '')
+
 
 def install_cache(
     cache_name: str = 'http_cache',
     backend: Optional[BackendSpecifier] = None,
+    module_only: bool = False,
     session_factory: Type[OriginalSession] = CachedSession,
     **kwargs,
 ):
@@ -33,22 +78,30 @@ def install_cache(
 
         >>> requests_cache.install_cache('demo_cache')
 
-    Accepts all the same parameters as :py:class:`.CachedSession`. Additional parameters:
+    Accepts all parameters for :py:class:`.CachedSession`. Additional parameters:
 
     Args:
+        module_only: Only install the cache for the current module
         session_factory: Session class to use. It must inherit from either
             :py:class:`.CachedSession` or :py:class:`.CacheMixin`
     """
     backend = init_backend(cache_name, backend, **kwargs)
+    module = _calling_module()
 
-    class _ConfiguredCachedSession(session_factory):  # type: ignore  # See mypy issue #5865
+    class _ConfiguredCachedSession(ModuleCacheMixin, session_factory):  # type: ignore  # See mypy issue #5865
         def __init__(self):
-            super().__init__(cache_name=cache_name, backend=backend, **kwargs)
+            super().__init__(
+                cache_name=cache_name,
+                backend=backend,
+                module_only=module_only,
+                modules=[module],
+                **kwargs,
+            )
 
     _patch_session_factory(_ConfiguredCachedSession)
 
 
-def uninstall_cache():
+def uninstall_cache(module_only: bool = False):
     """Disable the cache by restoring the original :py:class:`requests.Session`"""
     _patch_session_factory(OriginalSession)
 
@@ -96,6 +149,15 @@ def get_cache() -> Optional[BaseCache]:
     return getattr(requests.Session(), 'cache', None)
 
 
+def get_installed_modules() -> List[str]:
+    """Get all modules that have caching installed"""
+    session = requests.Session()
+    if isinstance(session, ModuleCacheMixin):
+        return session.modules
+    else:
+        return []
+
+
 def is_installed() -> bool:
     """Indicate whether or not requests-cache is currently installed"""
     return isinstance(requests.Session(), CachedSession)