Merge 03e092d into ae03f55

pynamodb/attributes.py

@@ -17,6 +17,7 @@
 from pynamodb.compat import getmembers_issubclass
 from pynamodb.expressions.operand import Path
 import collections
+import copy
 
 
 class Attribute(object):
@@ -386,6 +387,89 @@ def serialize(self, value):
             return six.u(value)
 
 
+class MultiKeyMap(object):
+    def __init__(self):
+        self.key = None
+        self.value = None
+        self.keymap = {}
+
+    def translate_key(self, key):
+        return self.keymap.get(key, key)
+
+    def add_keys(self, actual_key, keys):
+        self.keymap[actual_key]=actual_key
+        for key in keys:
+            self.keymap[key] = actual_key
+
+    def __getitem__(self, key):
+        if self.key != self.translate_key(key):
+            return None
+        return self.value
+
+    def __setitem__(self, key, value):
+        self.key = self.translate_key(key)
+        self.value = value
+
+    def __get__(self, instance, owner):
+        if instance:
+            return self.value
+        else:
+            return self
+
+    def get_key_value(self):
+        return (self.key, self.value)
+
+    def get_value(self):
+        return self.value
+
+class EnumAttribute(Attribute):
+    attr_type = STRING
+
+    def __init__(self, type_set={}, *args, **kwargs):
+        self.type_set = {ATTR_TYPE_MAP[x.attr_type]: x for x in type_set}
+        self.multikey_map = MultiKeyMap()
+        for t in type_set:
+            self.multikey_map.add_keys(ATTR_TYPE_MAP[t.attr_type], [t.attr_type, t.__class__, t.__class__.__name__ ])
+        Attribute.__init__(self, default=copy.deepcopy(self.multikey_map), *args, **kwargs)
+
+    def default(self):
+        return copy.deepcopy(self.multikey_map)
+
+    def serialize(self, value):
+        key, value = value.get_key_value()
+        if value is None:
+            return json.dumps(None)
+        return json.dumps({key: self.type_set[key].serialize(value)})
+
+    def deserialize(self, value):
+        value=json.loads(value)
+        ret = copy.deepcopy(self.multikey_map)
+        if value is None:
+            return ret
+        key = value.keys()[0]
+        value = value[key]
+
+        ret[key] = self.type_set[key].deserialize(value)
+        return ret
+
+    def __get__(self, instance, owner):
+         if instance:
+             attr_name = instance._dynamo_to_python_attrs.get(self.attr_name, self.attr_name)
+             return instance.attribute_values.get(attr_name, None)
+         else:
+             return self
+
+    def __eq__(self, other):
+        if other is None or isinstance(other, Attribute):  # handle object identity comparison
+            return self is other
+        key = other[0]
+        value = other[1]
+        ret = copy.deepcopy(self.multikey_map)
+        ret[key]=value
+        result = self.serialize(ret)
+        print (result)
+        return Path(self).__eq__(ret)
+
 class JSONAttribute(Attribute):
     """
     A JSON Attribute

showcase.py

@@ -0,0 +1,81 @@
+from pynamodb.attributes import (
+    UnicodeAttribute, UnicodeSetAttribute, Attribute, LegacyBooleanAttribute,
+    BooleanAttribute, NumberAttribute, BinaryAttribute,MapAttribute, ListAttribute, EnumAttribute)
+
+from pynamodb.models import Model
+
+class MyModel(Model):
+    class Meta:
+        table_name = "TestModel"
+        host="http://localhost:8000"
+    id = UnicodeAttribute(hash_key=True)
+    enum_attr = EnumAttribute(type_set={UnicodeAttribute(null=True), NumberAttribute(null=True), UnicodeSetAttribute(null=True)}, null=True)
+
+if MyModel.exists():
+    MyModel.delete_table()
+MyModel.create_table(read_capacity_units=2, write_capacity_units=2)
+
+
+############# Basic Usage ###############
+# Create an object
+m1 = MyModel("id1")
+m1.enum_attr['UnicodeAttribute'] = "Obama"
+m1.save()
+print (MyModel.dumps())
+# Outputs:
+# [["id1", {"attributes": {"enum_attr": {"S": "{\"S\": \"Obama\"}"}}}]]
+########################################
+
+# Read an object
+m1 = MyModel.get("id1")
+print (m1.enum_attr.get_value())
+# Outputs:
+# Obama
+########################################
+
+# Change value
+m1 = MyModel.get("id1")
+m1.enum_attr["UnicodeAttribute"] = "Obama"
+print (m1.enum_attr["UnicodeAttribute"])
+# Outputs:
+# Obama
+m1.enum_attr["NumberAttribute"] = 2018
+print (m1.enum_attr["NumberAttribute"])
+# Outputs:
+# 2018
+print (m1.enum_attr["UnicodeAttribute"])
+# Outputs:
+# None
+m1.save()
+print (MyModel.dumps())
+# Outputs:
+# [["id1", {"attributes": {"enum_attr": {"S": "{\"N\": \"2018\"}"}}}]]
+########################################
+
+# Different assignment patterns
+m2 = MyModel("id2")
+# All have the same effect
+m2.enum_attr["UnicodeAttribute"] = "President"
+m2.enum_attr["S"] = "President"
+m2.enum_attr["String"] = "President"
+m2.save()
+print (MyModel.dumps())
+# Outputs:
+# [["id1", {"attributes": {"enum_attr": {"S": "{\"N\": \"2018\"}"}}}], ["id2", {"attributes": {"enum_attr": {"S": "{\"S\": \"President\"}"}}}]] #noqa
+########################################
+
+# Nested values
+m3 = MyModel("id3")
+m3.enum_attr[UnicodeSetAttribute]={"Hi", "there", "World!"}
+m3.save()
+m3 = MyModel.get("id3")
+print (m3.enum_attr.get_value())
+# Outputs:
+#set([u'there', u'Hi', u'World!'])
+########################################
+
+# Querying
+for x in MyModel.query("id1", filter_condition=MyModel.enum_attr==(NumberAttribute,2018)):
+    print ("id: {}, enum_attr: {}".format(x.id, x.enum_attr.get_value()))
+# Outputs:
+#id: id1, enum_attr: 2018