from ..core import Machine, Transition, Event

from threading import RLock
import inspect


class LockedMethod:

    def __init__(self, lock, func):
        self.lock = lock
        self.func = func

    def __call__(self, *args, **kwargs):
        with self.lock:
            return self.func(*args, **kwargs)


class LockedEvent(Event):

    def trigger(self, model, *args, **kwargs):
        with self.machine.rlock:
            return super(LockedEvent, self).trigger(model, *args, **kwargs)


class LockedMachine(Machine):

    def __init__(self, *args, **kwargs):
        self.rlock = RLock()
        super(LockedMachine, self).__init__(*args, **kwargs)

    def __getattribute__(self, item):
        f = super(LockedMachine, self).__getattribute__
        tmp = f(item)
        if inspect.ismethod(tmp) and item not in "__getattribute__":
            return LockedMethod(f('rlock'), tmp)
        return tmp

    def __getattr__(self, item):
        try:
            return super(LockedMachine, self).__getattribute__(item)
        except AttributeError:
            return super(LockedMachine, self).__getattr__(item)

    @staticmethod
    def _create_event(*args, **kwargs):
        return LockedEvent(*args, **kwargs)

from ..core import Machine, Transition, Event, listify

from threading import RLock
import inspect

try:
    from contextlib import nested  # Python 2
except ImportError:
    from contextlib import ExitStack, contextmanager

    @contextmanager
    def nested(*contexts):
        """
        Reimplementation of nested in python 3.
        """
        with ExitStack() as stack:
            for ctx in contexts:
                stack.enter_context(ctx)
            yield contexts


class LockedMethod:

    def __init__(self, context, func):
        self.context = context
        self.func = func

    def __call__(self, *args, **kwargs):
        with nested(*self.context):
            return self.func(*args, **kwargs)


class LockedEvent(Event):

    def trigger(self, model, *args, **kwargs):
        with nested(*self.machine.context):
            return super(LockedEvent, self).trigger(model, *args, **kwargs)


class LockedMachine(Machine):

    def __init__(self, *args, **kwargs):
        try:
            self.context = listify(kwargs.pop('context'))
        except KeyError:
            self.context = [RLock()]

        super(LockedMachine, self).__init__(*args, **kwargs)

    def __getattribute__(self, item):
        f = super(LockedMachine, self).__getattribute__
        tmp = f(item)
        if inspect.ismethod(tmp) and item not in "__getattribute__":
            return LockedMethod(f('context'), tmp)
        return tmp

    def __getattr__(self, item):
        try:
            return super(LockedMachine, self).__getattribute__(item)
        except AttributeError:
            return super(LockedMachine, self).__getattr__(item)

    @staticmethod
    def _create_event(*args, **kwargs):
        return LockedEvent(*args, **kwargs)

Allow injecting a lock, or arbitrary context managers into LockedMachine
Allow injecting a lock, or arbitrary context managers into LockedMachine

- from ..core import Machine, Transition, Event
+ from ..core import Machine, Transition, Event, listify
  
  from threading import RLock
  import inspect
  
+ try:
+     from contextlib import nested  # Python 2
+ except ImportError:
+     from contextlib import ExitStack, contextmanager
+ 
+     @contextmanager
+     def nested(*contexts):
+         """
+         Reimplementation of nested in python 3.
+         """
+         with ExitStack() as stack:
+             for ctx in contexts:
+                 stack.enter_context(ctx)
+             yield contexts
+ 
  
  class LockedMethod:
  
-     def __init__(self, lock, func):
+     def __init__(self, context, func):
-         self.lock = lock
+         self.context = context
          self.func = func
  
      def __call__(self, *args, **kwargs):
-         with self.lock:
+         with nested(*self.context):
              return self.func(*args, **kwargs)
  
  
  class LockedEvent(Event):
  
      def trigger(self, model, *args, **kwargs):
-         with self.machine.rlock:
+         with nested(*self.machine.context):
              return super(LockedEvent, self).trigger(model, *args, **kwargs)
  
  
  class LockedMachine(Machine):
  
      def __init__(self, *args, **kwargs):
+         try:
+             self.context = listify(kwargs.pop('context'))
+         except KeyError:
-         self.rlock = RLock()
+             self.context = [RLock()]
+ 
          super(LockedMachine, self).__init__(*args, **kwargs)
  
      def __getattribute__(self, item):
          f = super(LockedMachine, self).__getattribute__
          tmp = f(item)
          if inspect.ismethod(tmp) and item not in "__getattribute__":
-             return LockedMethod(f('rlock'), tmp)
+             return LockedMethod(f('context'), tmp)
          return tmp
  
      def __getattr__(self, item):
          try:
              return super(LockedMachine, self).__getattribute__(item)
          except AttributeError:
              return super(LockedMachine, self).__getattr__(item)
  
      @staticmethod
      def _create_event(*args, **kwargs):
          return LockedEvent(*args, **kwargs)
  
Allow injecting a lock, or arbitrary context managers into LockedMachine
## Code Before:
from ..core import Machine, Transition, Event

from threading import RLock
import inspect


class LockedMethod:

    def __init__(self, lock, func):
        self.lock = lock
        self.func = func

    def __call__(self, *args, **kwargs):
        with self.lock:
            return self.func(*args, **kwargs)


class LockedEvent(Event):

    def trigger(self, model, *args, **kwargs):
        with self.machine.rlock:
            return super(LockedEvent, self).trigger(model, *args, **kwargs)


class LockedMachine(Machine):

    def __init__(self, *args, **kwargs):
        self.rlock = RLock()
        super(LockedMachine, self).__init__(*args, **kwargs)

    def __getattribute__(self, item):
        f = super(LockedMachine, self).__getattribute__
        tmp = f(item)
        if inspect.ismethod(tmp) and item not in "__getattribute__":
            return LockedMethod(f('rlock'), tmp)
        return tmp

    def __getattr__(self, item):
        try:
            return super(LockedMachine, self).__getattribute__(item)
        except AttributeError:
            return super(LockedMachine, self).__getattr__(item)

    @staticmethod
    def _create_event(*args, **kwargs):
        return LockedEvent(*args, **kwargs)

## Instruction:
Allow injecting a lock, or arbitrary context managers into LockedMachine
## Code After:
from ..core import Machine, Transition, Event, listify

from threading import RLock
import inspect

try:
    from contextlib import nested  # Python 2
except ImportError:
    from contextlib import ExitStack, contextmanager

    @contextmanager
    def nested(*contexts):
        """
        Reimplementation of nested in python 3.
        """
        with ExitStack() as stack:
            for ctx in contexts:
                stack.enter_context(ctx)
            yield contexts


class LockedMethod:

    def __init__(self, context, func):
        self.context = context
        self.func = func

    def __call__(self, *args, **kwargs):
        with nested(*self.context):
            return self.func(*args, **kwargs)


class LockedEvent(Event):

    def trigger(self, model, *args, **kwargs):
        with nested(*self.machine.context):
            return super(LockedEvent, self).trigger(model, *args, **kwargs)


class LockedMachine(Machine):

    def __init__(self, *args, **kwargs):
        try:
            self.context = listify(kwargs.pop('context'))
        except KeyError:
            self.context = [RLock()]

        super(LockedMachine, self).__init__(*args, **kwargs)

    def __getattribute__(self, item):
        f = super(LockedMachine, self).__getattribute__
        tmp = f(item)
        if inspect.ismethod(tmp) and item not in "__getattribute__":
            return LockedMethod(f('context'), tmp)
        return tmp

    def __getattr__(self, item):
        try:
            return super(LockedMachine, self).__getattribute__(item)
        except AttributeError:
            return super(LockedMachine, self).__getattr__(item)

    @staticmethod
    def _create_event(*args, **kwargs):
        return LockedEvent(*args, **kwargs)
