Merged in jaraco/six (pull request #12)

Add 'patch_with_metaclass' which provides a more complete interface for supporting metaclass-based classes across Python versions.

3 files changed, 131 insertions, 0 deletions

diff --git a/documentation/index.rst b/documentation/index.rst
index 473e39b..6d43163 100644
--- a/documentation/index.rst
+++ b/documentation/index.rst
@@ -280,6 +280,36 @@ Python 2 and 3.
           pass
 
 
+.. function:: add_metaclass(metaclass)
+
+   Decorate a class to replace it with a metaclass-constructed version.
+   Similar to :func:`with_metaclass`, but allows the resultant class to be
+   created without an intermediate class. For example::
+
+       @add_metaclass(MyMeta)
+       class MyClass(object):
+           ...
+
+   That code produces a class equivalent to ::
+
+       class MyClass(object, metaclass=MyMeta):
+           ...
+
+   on Python 3 or ::
+
+       class MyClass(object):
+           __metaclass__ = MyMeta
+
+   on Python 2.
+
+   Requires Python 2.6 or later (for class decoration). For use on Python
+   2.5 and earlier, use the legacy syntax::
+
+       class MyClass(object):
+           ...
+       MyClass = add_metaclass(MyClass)
+
+
 Binary and text data
 >>>>>>>>>>>>>>>>>>>>
 
diff --git a/six.py b/six.py
index f9af45c..df5e537 100644
--- a/six.py
+++ b/six.py
@@ -423,3 +423,41 @@ _add_doc(reraise, """Reraise an exception.""")
 def with_metaclass(meta, *bases):
     """Create a base class with a metaclass."""
     return meta("NewBase", bases, {})
+
+def add_metaclass(metaclass):
+    """
+    Decorate a class to replace it with a metaclass-constructed version.
+
+    Usage:
+
+    @add_metaclass(MyMeta)
+    class MyClass(object):
+        ...
+
+    That code produces a class equivalent to
+
+    class MyClass(object, metaclass=MyMeta):
+        ...
+
+    on Python 3 or
+
+    class MyClass(object):
+        __metaclass__ = MyMeta
+
+    on Python 2
+
+    Requires Python 2.6 or later (for class decoration). For use on Python
+    2.5 and earlier, use the legacy syntax:
+
+    class MyClass(object):
+        ...
+    MyClass = add_metaclass(MyClass)
+    """
+    def wrapper(cls):
+        orig_vars = cls.__dict__.copy()
+        orig_vars.pop('__dict__', None)
+        orig_vars.pop('__weakref__', None)
+        for slots_var in orig_vars.get('__slots__', ()):
+            orig_vars.pop(slots_var)
+        return metaclass(cls.__name__, cls.__bases__, orig_vars)
+    return wrapper
diff --git a/test_six.py b/test_six.py
index 38bd0e5..12d3a06 100644
--- a/test_six.py
+++ b/test_six.py
@@ -511,3 +511,66 @@ def test_with_metaclass():
     assert type(X) is Meta
     assert issubclass(X, Base)
     assert issubclass(X, Base2)
+
+
+def test_add_metaclass():
+    class Meta(type):
+        pass
+    class X:
+        "success"
+    X = six.add_metaclass(Meta)(X)
+    assert type(X) is Meta
+    assert issubclass(X, object)
+    assert X.__module__ is __name__
+    assert X.__doc__ == "success"
+    class Base(object):
+        pass
+    class X(Base):
+        pass
+    X = six.add_metaclass(Meta)(X)
+    assert type(X) is Meta
+    assert issubclass(X, Base)
+    class Base2(object):
+        pass
+    class X(Base, Base2):
+        pass
+    X = six.add_metaclass(Meta)(X)
+    assert type(X) is Meta
+    assert issubclass(X, Base)
+    assert issubclass(X, Base2)
+
+    # additionally test a second-generation subclass of a type
+    class Meta1(type):
+        m1 = "m1"
+    class Meta2(Meta1):
+        m2 = "m2"
+    class Base:
+        b = "b"
+    Base = six.add_metaclass(Meta1)(Base)
+    class X(Base):
+        x = "x"
+    X = six.add_metaclass(Meta2)(X)
+    assert type(X) is Meta2
+    assert issubclass(X, Base)
+    assert type(Base) is Meta1
+    assert "__dict__" not in vars(X)
+    instance = X()
+    instance.attr = "test"
+    assert vars(instance) == {"attr": "test"}
+    assert instance.b == Base.b
+    assert instance.x == X.x
+
+    # test a class with slots
+    class MySlots(object):
+        __slots__ = ["a", "b"]
+    MySlots = six.add_metaclass(Meta1)(MySlots)
+
+    assert MySlots.__slots__ == ["a", "b"]
+    instance = MySlots()
+    instance.a = "foo"
+    try:
+        instance.c = "baz"
+    except AttributeError:
+        pass
+    else:
+        raise RuntimeError("did not raise")