from typing import Any, 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"):
        self.obj = obj
        self.dumps = getattr(obj, dumps)
        self.loads = getattr(obj, 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)]

    def dumps(self, value) -> Union[str, bytes]:
        for step in self.dump_steps:
            value = step(value)
        return value

    def loads(self, value) -> CachedResponse:
        for step in self.load_steps:
            value = step(value)
        return value

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

    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) 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.
    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:
            value = step(value)
        return value

    def loads(self, value) -> CachedResponse:
        for step in self.load_steps:
            value = step(value)
        return value

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

- 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"""
+     """Generic class to wrap serialization steps with consistent ``dumps()`` and ``loads()`` methods
  
-     def __init__(self, obj: Any, dumps: str = "dumps", loads: str = "loads"):
+     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"""
+     """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, steps: List):
+     def __init__(self, stages: List):
-         self.steps = steps
+         self.steps = stages
-         self.dump_steps = [step.dumps for step in steps]
+         self.dump_steps = [step.dumps for step in stages]
-         self.load_steps = [step.loads for step in reversed(steps)]
+         self.load_steps = [step.loads for step in reversed(stages)]
  
      def dumps(self, value) -> Union[str, bytes]:
          for step in self.dump_steps:
              value = step(value)
          return value
  
      def loads(self, value) -> CachedResponse:
          for step in self.load_steps:
              value = step(value)
          return value
  
Allow Stage objects to take functions instead of object + method names
## Code Before:
from typing import Any, 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"):
        self.obj = obj
        self.dumps = getattr(obj, dumps)
        self.loads = getattr(obj, 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)]

    def dumps(self, value) -> Union[str, bytes]:
        for step in self.dump_steps:
            value = step(value)
        return value

    def loads(self, value) -> CachedResponse:
        for step in self.load_steps:
            value = step(value)
        return value

## Instruction:
Allow Stage objects to take functions instead of object + method names
## Code After:
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

    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) 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.
    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:
            value = step(value)
        return value

    def loads(self, value) -> CachedResponse:
        for step in self.load_steps:
            value = step(value)
        return value
