WIP: shenanigans

1 files changed, 144 insertions, 72 deletions

diff --git a/requests_cache/patcher.py b/requests_cache/patcher.py
index 36143d6..ac58be2 100644
--- a/requests_cache/patcher.py
+++ b/requests_cache/patcher.py
@@ -18,6 +18,13 @@ import requests
 from .backends import BackendSpecifier, BaseCache, init_backend
 from .session import CachedSession, OriginalSession
 
+# TODO: This is going to require thinking through many more edge cases.
+#   Lots of ways this could potentially go wrong.
+# TODO: What if we want to use module_only but patch an external library? Options:
+#   1. Check if enabled module is anywhere in the call stack (e.g., it called some other module that
+#      made the request)
+#   2. Provide target module to patch (e.g., 'github.Requester')
+
 logger = getLogger(__name__)
 
 if TYPE_CHECKING:
@@ -26,67 +33,11 @@ else:
     MIXIN_BASE = object
 
 
-# TODO: This is going to require thinking through many more edge cases. Lots of ways this could go wrong.
-class ModuleCacheMixin(MIXIN_BASE):
-    """Session mixin that only caches requests sent from specific modules. May be used in one of two
-    modes:
-
-    * Opt-in: caching is disabled by default, and enabled for modules in ``include_modules``
-    * Opt-out: caching is enabled by default, and disabled for modules in ``exclude_modules``
-
-    Args:
-        include_modules: List of modules to enable caching for
-        exclude_modules: List of modules to disable caching for
-        opt_in: Whether to use opt-in mode (``True``) or opt-out mode (``False``)
-    """
-
-    def __init__(
-        self,
-        *args,
-        include_modules: Optional[List[str]] = None,
-        exclude_modules: Optional[List[str]] = None,
-        opt_in: bool = True,
-        **kwargs,
-    ):
-        super().__init__(*args, **kwargs)
-        self.include_modules = set(include_modules or [])
-        self.exclude_modules = set(exclude_modules or [])
-        self.opt_in = opt_in
-
-    def request(self, *args, **kwargs):
-        if self.is_module_enabled(back=3):
-            return super().request(*args, **kwargs)
-        else:
-            return OriginalSession.request(self, *args, **kwargs)
-
-    def is_module_enabled(self, back: int = 2) -> bool:
-        module = _calling_module(back=back)
-        if self.opt_in:
-            return module in self.include_modules
-        else:
-            return module not in self.exclude_modules
-
-    def enable_module(self):
-        module = _calling_module()
-        if self.opt_in:
-            self.include_modules |= {module}
-        else:
-            self.exclude_modules -= {module}
-        logger.info(f'Caching enabled for {module}')
-
-    def disable_module(self):
-        module = _calling_module()
-        if self.opt_in:
-            self.include_modules -= {module}
-        else:
-            self.exclude_modules |= {module}
-        logger.info(f'Caching disabled for {module}')
-
-
 def install_cache(
     cache_name: str = 'http_cache',
     backend: Optional[BackendSpecifier] = None,
     module_only: bool = False,
+    session: CachedSession = None,
     session_factory: Type[OriginalSession] = CachedSession,
     **kwargs,
 ):
@@ -101,19 +52,40 @@ def install_cache(
 
     Args:
         module_only: Only install the cache for the current module
+        session: An existing session object to use
         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)
-    if module_only:
+    # Patch with an existing session object
+    if session:
+        cls = _get_session_wrapper(session)
+        _patch_session_factory(cls)
+    # Patch only for the current module
+    elif module_only:
         modules = get_installed_modules() + [_calling_module()]
-        _install_modules(cache_name, backend, session_factory, modules, **kwargs)
-
-    class _ConfiguredCachedSession(session_factory):  # type: ignore  # See mypy issue #5865
-        def __init__(self):
-            super().__init__(cache_name=cache_name, backend=backend, **kwargs)
-
-    _patch_session_factory(_ConfiguredCachedSession)
+        # _install_modules(
+        #     cache_name,
+        #     init_backend(cache_name, backend, **kwargs),
+        #     modules,
+        #     **kwargs,
+        # )
+        cls = _get_configured_session(
+            ModuleCachedSession,
+            cache_name=cache_name,
+            backend=backend,
+            include_modules=modules,
+            **kwargs,
+        )
+        _patch_session_factory(cls)
+    # Patch with CachedSession or session_factory
+    else:
+        cls = _get_configured_session(
+            session_factory,
+            cache_name=cache_name,
+            backend=init_backend(cache_name, backend, **kwargs),
+            **kwargs,
+        )
+        _patch_session_factory(cls)
 
 
 def uninstall_cache(module_only: bool = False):
@@ -174,7 +146,7 @@ def get_cache() -> Optional[BaseCache]:
 def get_installed_modules() -> List[str]:
     """Get all modules that have caching installed"""
     session = requests.Session()
-    if isinstance(session, ModuleCacheMixin):
+    if isinstance(session, ModuleCachedSession):
         return list(session.include_modules)
     else:
         return []
@@ -183,7 +155,7 @@ def get_installed_modules() -> List[str]:
 def is_installed() -> bool:
     """Indicate whether or not requests-cache is currently installed"""
     session = requests.Session()
-    if isinstance(session, ModuleCacheMixin):
+    if isinstance(session, ModuleCachedSession):
         return session.is_module_enabled()
     else:
         return isinstance(session, CachedSession)
@@ -213,6 +185,71 @@ def remove_expired_responses():
     delete(expired=True)
 
 
+class ModuleCachedSession(CachedSession):
+    """Session mixin that only caches requests sent from specific modules. May be used in one of two
+    modes:
+
+    * Opt-in: caching is disabled by default, and enabled for modules in ``include_modules``
+    * Opt-out: caching is enabled by default, and disabled for modules in ``exclude_modules``
+
+    Args:
+        include_modules: List of modules to enable caching for
+        exclude_modules: List of modules to disable caching for
+        opt_in: Whether to use opt-in mode (``True``) or opt-out mode (``False``)
+    """
+
+    def __init__(
+        self,
+        *args,
+        include_modules: List[str] = None,
+        exclude_modules: List[str] = None,
+        opt_in: bool = True,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.include_modules = set(include_modules or [])
+        self.exclude_modules = set(exclude_modules or [])
+        self.opt_in = opt_in
+
+    def close(self):
+        """Only close adapter(s) and skip closing backend connections in CachedSession.close(), as
+        the object may be reused after closing. :py:func:`requests.request` will create a session
+        for the request and close it after use. We will keep this behavior to avoid memory leaks.
+        Adapter(s) will create new connections if re-used after closing.
+        """
+        for v in self.adapters.values():
+            v.close()
+
+    def disable_module(self):
+        module = _calling_module()
+        if self.opt_in:
+            self.include_modules -= {module}
+        else:
+            self.exclude_modules |= {module}
+        logger.info(f'Caching disabled for {module}')
+
+    def enable_module(self):
+        module = _calling_module()
+        if self.opt_in:
+            self.include_modules |= {module}
+        else:
+            self.exclude_modules -= {module}
+        logger.info(f'Caching enabled for {module}')
+
+    def is_module_enabled(self, back: int = 2) -> bool:
+        module = _calling_module(back=back)
+        if self.opt_in:
+            return module in self.include_modules
+        else:
+            return module not in self.exclude_modules
+
+    def request(self, *args, **kwargs):
+        if self.is_module_enabled(back=3):
+            return super().request(*args, **kwargs)
+        else:
+            return OriginalSession.request(self, *args, **kwargs)
+
+
 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
@@ -220,16 +257,31 @@ def _calling_module(back: int = 2) -> str:
     return getattr(module, '__name__', '')
 
 
+# testing
+def _calling_module_2() -> str:
+    """Get the name of the calling module (first module outside of requests-cache)"""
+    for module in _stack_modules():
+        if not module.startswith('requests_cache'):
+            return module
+
+    raise RuntimeError('Could not determine calling module')
+
+
+def _stack_modules() -> Iterator[str]:
+    for frame_info in inspect.stack():
+        module = inspect.getmodule(frame_info.frame)
+        yield getattr(module, '__name__', '')
+
+
 def _install_modules(
     cache_name: str,
     backend: BackendSpecifier,
-    session_factory: Type[OriginalSession],
     modules: List[str],
     **kwargs,
 ):
     """Install the cache for specific modules"""
 
-    class _ConfiguredCachedSession(ModuleCacheMixin, session_factory):  # type: ignore  # See mypy issue #5865
+    class _ConfiguredCachedSession(ModuleCachedSession):  # type: ignore  # See mypy issue #5865
         def __init__(self):
             super().__init__(
                 cache_name=cache_name, backend=backend, include_modules=modules, **kwargs
@@ -241,7 +293,7 @@ def _install_modules(
 def _uninstall_module():
     """Uninstall the cache for the current module"""
     session = requests.Session()
-    if not isinstance(session, ModuleCacheMixin):
+    if not isinstance(session, ModuleCachedSession):
         return
 
     modules = get_installed_modules()
@@ -261,6 +313,26 @@ def _uninstall_module():
         )
 
 
-def _patch_session_factory(session_factory: Type[OriginalSession] = CachedSession):
+def _get_configured_session(
+    session_factory: Type[OriginalSession], *args, **kwargs
+) -> Type[OriginalSession]:
+    class _ConfiguredCachedSession(session_factory):  # type: ignore  # See mypy issue #5865
+        def __init__(self):
+            super().__init__(*args, **kwargs)
+
+    return _ConfiguredCachedSession
+
+
+def _get_session_wrapper(session: CachedSession):
+    """Create a wrapper Session class that returns an existing session object when created"""
+
+    class _SessionWrapper(OriginalSession):
+        def __new__(self, *args, **kwargs):
+            return session
+
+    return _SessionWrapper
+
+
+def _patch_session_factory(session_factory: Type[OriginalSession]):
     logger.debug(f'Patching requests.Session with class: {session_factory.__name__}')
     requests.Session = requests.sessions.Session = session_factory  # type: ignore