Allow Stage objects to take functions instead of object + method names

1 files changed, 26 insertions, 12 deletions

diff --git a/requests_cache/serializers/pipeline.py b/requests_cache/serializers/pipeline.py
index fec7889..3bbfbda 100644
--- a/requests_cache/serializers/pipeline.py
+++ b/requests_cache/serializers/pipeline.py
@@ -3,27 +3,41 @@
    :classes-only:
    :nosignatures:
 """
-from typing import Any, List, Union
+from typing import Any, Callable, List, Union
 
 from ..models import CachedResponse
 
 
 class Stage:
-    """Generic class to wrap serialization steps with consistent ``dumps()`` and ``loads()`` methods"""
-
-    def __init__(self, obj: Any, dumps: str = "dumps", loads: str = "loads"):
+    """Generic class to wrap serialization steps with consistent ``dumps()`` and ``loads()`` methods
+
+    Args:
+        obj: Serializer object or module, if applicable
+        dumps: Serialization function, or name of method on ``obj``
+        loads: Deserialization function, or name of method on ``obj``
+    """
+
+    def __init__(
+        self,
+        obj: Any = None,
+        dumps: Union[str, Callable] = 'dumps',
+        loads: Union[str, Callable] = 'loads',
+    ):
         self.obj = obj
-        self.dumps = getattr(obj, dumps)
-        self.loads = getattr(obj, loads)
+        self.dumps = getattr(obj, dumps) if isinstance(dumps, str) else dumps
+        self.loads = getattr(obj, loads) if isinstance(loads, str) else loads
 
 
 class SerializerPipeline:
-    """A sequence of steps used to serialize and deserialize response objects"""
-
-    def __init__(self, steps: List):
-        self.steps = steps
-        self.dump_steps = [step.dumps for step in steps]
-        self.load_steps = [step.loads for step in reversed(steps)]
+    """A sequence of steps used to serialize and deserialize response objects.
+    This can be initialized with :py:class:`Stage` objects, or any objects with ``dumps()`` and
+    ``loads()`` methods
+    """
+
+    def __init__(self, stages: List):
+        self.steps = stages
+        self.dump_steps = [step.dumps for step in stages]
+        self.load_steps = [step.loads for step in reversed(stages)]
 
     def dumps(self, value) -> Union[str, bytes]:
         for step in self.dump_steps: