Add support for defer & recover

`defer` allows to schedule a cleanup to be executed when current function
returns. It is similar to `try`/`finally` but does not force the cleanup part
to be far away in the end. For example::

   wc = wcfs.join(zurl)    │     wc = wcfs.join(zurl)
   defer(wc.close)         │     try:
                           │        ...
   ...                     │        ...
   ...                     │        ...
   ...                     │     finally:
                           │        wc.close()

For completeness there is `recover` and `panic` that allow to program with
Go-style error handling, for example::

   def _():
      r = recover()
      if r is not None:
         print("recovered. error was: %s" % (r,))
   defer(_)

   ...

   panic("aaa")

But `recover` and `panic` are probably of less utility since they can be
practically natively modelled with `try`/`except`.

If `defer` is used, the function that uses it must be wrapped with `@func` or
`@method` decorators.

The implementation is partly inspired by work of Denis Kolodin:

- https://habr.com/post/191786
- https://stackoverflow.com/a/43028386/9456786

README.rst

@@ -7,6 +7,7 @@ Package golang provides Go-like features for Python:
 - `go` spawns lightweight thread.
 - `chan` and `select` provide channels with Go semantic.
 - `method` allows to define methods separate from class.
+- `defer` allows to schedule a cleanup from the main control flow.
 - `gimport` allows to import python modules by full path in a Go workspace.
 
 
@@ -63,6 +64,40 @@ For example::
 will define `MyClass.my_method()`.
 
 
+Defer / recover / panic
+-----------------------
+
+`defer` allows to schedule a cleanup to be executed when current function
+returns. It is similar to `try`/`finally` but does not force the cleanup part
+to be far away in the end. For example::
+
+   wc = wcfs.join(zurl)    │     wc = wcfs.join(zurl)
+   defer(wc.close)         │     try:
+                           │        ...
+   ...                     │        ...
+   ...                     │        ...
+   ...                     │     finally:
+                           │        wc.close()
+
+For completeness there is `recover` and `panic` that allow to program with
+Go-style error handling, for example::
+
+   def _():
+      r = recover()
+      if r is not None:
+         print("recovered. error was: %s" % (r,))
+   defer(_)
+
+   ...
+
+   panic("aaa")
+
+But `recover` and `panic` are probably of less utility since they can be
+practically natively modelled with `try`/`except`.
+
+If `defer` is used, the function that uses it must be wrapped with `@func` or
+`@method` decorators.
+
 Import
 ------
 

golang/__init__.py

@@ -22,15 +22,17 @@
 - `go` spawns lightweight thread.
 - `chan` and `select` provide channels with Go semantic.
 - `method` allows to define methods separate from class.
+- `defer` allows to schedule a cleanup from the main control flow.
 - `gimport` allows to import python modules by full path in a Go workspace.
 
 ...
 """
 
-__all__ = ['method', 'go', 'chan', 'select', 'default', 'panic', 'gimport']
+__all__ = ['method', 'go', 'chan', 'select', 'default', 'defer', 'panic', 'recover', 'func' 'gimport']
 
 from golang._gopath import gimport  # make gimport available from golang
-import threading, collections, random
+import inspect, threading, collections, random, sys
+import decorator
 
 import six
 from golang._pycompat import im_class
@@ -58,6 +60,9 @@ from golang._pycompat import im_class
 #       ...
 def method(cls):
     def deco(f):
+        # wrap f with @func, so that e.g. defer works automatically.
+        f = func(f)
+
         if isinstance(f, (staticmethod, classmethod)):
             func_name = f.__func__.__name__
         else:
@@ -73,6 +78,96 @@ def panic(arg):
 class _PanicError(Exception):
     pass
 
+
+# @func is a necessary decorator for functions for selected golang features to work.
+#
+# It is needed for defer.
+def func(f):
+    # @staticmethod & friends require special care:
+    # unpack f first to original func and then repack back after wrapping.
+    fclass = None
+    if isinstance(f, (staticmethod, classmethod)):
+        fclass = type(f)
+        f = f.__func__
+
+    def _(f, *argv, **kw):
+        # run f under separate frame, where defer will register calls.
+        __goframe__ = _GoFrame()
+        with __goframe__:
+            return f(*argv, **kw)
+
+    # keep all f attributes, like __name__, __doc__, etc on _
+    _ = decorator.decorate(f, _)
+
+    # repack _ into e.g. @staticmethod if that was used on f.
+    if fclass is not None:
+        _ = fclass(_)
+
+    return _
+
+# _GoFrame serves __goframe__ that is setup by @func.
+class _GoFrame:
+    def __init__(self):
+        self.deferv    = []     # defer registers funcs here
+        self.recovered = False  # whether exception, if there was any, was recovered
+
+    def __enter__(self):
+        pass
+
+    # __exit__ simulates both except and finally.
+    def __exit__(__goframe__, exc_type, exc_val, exc_tb):
+        if exc_val is not None:
+            __goframe__.recovered = False
+
+        if len(__goframe__.deferv) != 0:
+            d = __goframe__.deferv.pop()
+
+            # even if d panics - we have to call other defers
+            with __goframe__:
+                d()
+
+        return __goframe__.recovered
+
+# recover checks whether there is exception/panic currently being raised and returns it.
+#
+# If it was panic - it returns the argument that was passed to panic.
+# If there is other exception - it returns the exception object.
+#
+# If there is no exception/panic, or the panic argument was None - recover returns None.
+# Recover also returns None if it was not called by a deferred function directly.
+def recover():
+    fcall = inspect.currentframe().f_back   # caller's frame (deferred func)
+    fgo   = fcall.f_back                    # caller's parent frame defined by _GoFrame.__exit__
+    try:
+        goframe = fgo.f_locals['__goframe__']
+    except KeyError:
+        # called not under go func/defer
+        return None
+
+    _, exc, _ = sys.exc_info()
+    if exc is not None:
+        goframe.recovered = True
+    if type(exc) is _PanicError:
+        exc = exc.args[0]
+    return exc
+
+# defer registers f to be called when caller function exits.
+#
+# It is similar to try/finally but does not force the cleanup part to be far
+# away in the end.
+def defer(f):
+    fcall = inspect.currentframe().f_back   # caller's frame
+    fgo   = fcall.f_back                    # caller's parent frame defined by @func
+    try:
+        goframe = fgo.f_locals['__goframe__']
+    except KeyError:
+        panic("function %s uses defer, but not @func" % fcall.f_code.co_name)
+
+    goframe.deferv.append(f)
+
+
+
+
 # go spawns lightweight thread.
 #
 # NOTE it spawns threading.Thread, but if gevent was activated via

golang/golang_test.py

@@ -18,7 +18,7 @@
 # See COPYING file for full licensing terms.
 # See https://www.nexedi.com/licensing for rationale and options.
 
-from golang import go, chan, select, default, _PanicError, method
+from golang import go, chan, select, default, _PanicError, func, method, panic, defer, recover
 from pytest import raises
 from os.path import dirname
 import os, sys, time, threading, inspect, subprocess
@@ -385,6 +385,7 @@ def test_select():
 
 def test_method():
     # test how @method works
+    # this also implicitly tests @func, since @method uses that.
 
     class MyClass:
         def __init__(self, v):
@@ -413,7 +414,7 @@ def test_method():
     assert obj.mstatic(5)   == 5 + 1
     assert obj.mcls(7)      == 7 + 1
 
-    # this tests that @method preserves decorated function signature
+    # this tests that @func (used by @method) preserves decorated function signature
     assert inspect.formatargspec(*inspect.getargspec(MyClass.zzz)) == '(self, v, x=2, **kkkkwww)'
 
     assert MyClass.zzz.__module__       == __name__
@@ -424,3 +425,205 @@ def test_method():
 
     assert MyClass.mcls.__module__      == __name__
     assert MyClass.mcls.__name__        == 'mcls'
+
+
+
+def test_deferrecover():
+    # regular defer calls
+    v = []
+    @func
+    def _():
+        defer(lambda: v.append(1))
+        defer(lambda: v.append(2))
+        defer(lambda: v.append(3))
+
+    _()
+    assert v == [3, 2, 1]
+
+    # defers called even if exception is raised
+    v = []
+    @func
+    def _():
+        defer(lambda: v.append(1))
+        defer(lambda: v.append(2))
+        def _(): v.append('ran ok')
+        defer(_)
+        1/0
+
+    raises(ZeroDivisionError, "_()")
+    assert v == ['ran ok', 2, 1]
+
+    # defer without @func is caught and properly reported
+    v = []
+    def nofunc():
+        defer(lambda: v.append('xx'))
+
+    with raises(_PanicError) as exc:
+        nofunc()
+    assert exc.value.args == ("function nofunc uses defer, but not @func",)
+
+
+    # panic in deferred call - all defers are called
+    v = []
+    @func
+    def _():
+        defer(lambda: v.append(1))
+        defer(lambda: v.append(2))
+        defer(lambda: panic(3))
+        defer(lambda: v.append(4))
+
+    raises(_PanicError, "_()")
+    assert v == [4, 2, 1]
+
+
+    # defer + recover
+    v = []
+    @func
+    def _():
+        defer(lambda: v.append(1))
+        def _():
+            r = recover()
+            assert r == "aaa"
+            v.append('recovered ok')
+        defer(_)
+        defer(lambda: v.append(3))
+
+        panic("aaa")
+
+    _()
+    assert v == [3, 'recovered ok', 1]
+
+
+    # recover + panic in defer
+    v = []
+    @func
+    def _():
+        defer(lambda: v.append(1))
+        defer(lambda: panic(2))
+        def _():
+            r = recover()
+            assert r == "bbb"
+            v.append('recovered 1')
+        defer(_)
+        defer(lambda: v.append(3))
+
+        panic("bbb")
+
+    raises(_PanicError, "_()")
+    assert v == [3, 'recovered 1', 1]
+
+
+    # recover + panic in defer + recover
+    v = []
+    @func
+    def _():
+        defer(lambda: v.append(1))
+        def _():
+            r = recover()
+            assert r == "ddd"
+            v.append('recovered 2')
+        defer(_)
+        defer(lambda: panic("ddd"))
+        def _():
+            r = recover()
+            assert r == "ccc"
+            v.append('recovered 1')
+        defer(_)
+        defer(lambda: v.append(3))
+
+        panic("ccc")
+
+    _()
+    assert v == [3, 'recovered 1', 'recovered 2', 1]
+
+
+    # ---- recover() -> None ----
+
+    # no exception / not under defer
+    assert recover() is None
+
+    # no exception/panic
+    @func
+    def _():
+        def _():
+            assert recover() is None
+        defer(_)
+
+    # not directly called by deferred func
+    v = []
+    @func
+    def _():
+        def f():
+            assert recover() is None
+            v.append('not recovered')
+        defer(lambda: f())
+
+        panic("zzz")
+
+    raises(_PanicError, "_()")
+    assert v == ['not recovered']
+
+
+    # ---- defer in @method(x) ----
+
+    # defer in @method
+    v = []
+
+    class MyClass:
+        pass
+
+    @method(MyClass)
+    def zzz(self):
+        defer(lambda: v.append(1))
+        defer(lambda: v.append(2))
+        defer(lambda: v.append(3))
+
+    obj = MyClass()
+    obj.zzz()
+    assert v == [3, 2, 1]
+
+
+    # defer in std method
+    v = []
+
+    class MyClass:
+        @func
+        def method(self):
+            defer(lambda: v.append(1))
+            defer(lambda: v.append(2))
+            defer(lambda: v.append(4))
+
+    obj = MyClass()
+    obj.method()
+    assert v == [4, 2, 1]
+
+
+    # defer in std @staticmethod
+    v = []
+
+    class MyClass:
+        @func
+        @staticmethod
+        def mstatic():
+            defer(lambda: v.append(1))
+            defer(lambda: v.append(2))
+            defer(lambda: v.append(5))
+
+    MyClass.mstatic()
+    assert v == [5, 2, 1]
+
+
+    # defer in std @classmethod
+    v = []
+
+    class MyClass:
+        @func
+        @classmethod
+        def mcls(cls):
+            assert cls is MyClass
+            defer(lambda: v.append(1))
+            defer(lambda: v.append(2))
+            defer(lambda: v.append(7))
+
+    MyClass.mcls()
+    assert v == [7, 2, 1]

setup.py

@@ -21,7 +21,7 @@ setup(
 
     packages    = find_packages(),
 
-    install_requires = ['six'],
+    install_requires = ['six', 'decorator'],
 
     extras_require = {
                   'test': ['pytest'],