Merge pull request #639 from dimitri-yatsenko/custom-attribute-type

Custom attribute type: implement #627

Author: guzman-raphael

datajoint/__init__.py

@@ -23,7 +23,7 @@
            'Manual', 'Lookup', 'Imported', 'Computed', 'Part',
            'Not', 'AndList', 'U', 'Diagram', 'Di', 'ERD',
            'set_password', 'kill',
-           'MatCell', 'MatStruct',
+           'MatCell', 'MatStruct', 'AttributeAdapter',
            'errors', 'DataJointError', 'key']
 
 from .version import __version__
@@ -38,6 +38,7 @@
 from .admin import set_password, kill
 from .blob import MatCell, MatStruct
 from .fetch import key
+from .attribute_adapter import AttributeAdapter
 from . import errors
 from .errors import DataJointError
 

datajoint/attribute_adapter.py

@@ -0,0 +1,73 @@
+import re
+import os
+from .errors import DataJointError
+
+
+ADAPTED_TYPE_SWITCH = "DJ_SUPPORT_ADAPTED_TYPES"
+
+
+def _switch_adapated_types(on):
+    """
+    Enable (on=True) or disable (on=False) support for AttributeAdapter
+    """
+    if on:
+        os.environ[ADAPTED_TYPE_SWITCH] = "TRUE"
+    else:
+        del os.environ[ADAPTED_TYPE_SWITCH]
+
+
+def _support_adapted_types():
+    """
+    check if support for AttributeAdapter is enabled
+    """
+    return os.getenv(ADAPTED_TYPE_SWITCH, "FALSE").upper() == "TRUE"
+
+
+class AttributeAdapter:
+    """
+    Base class for adapter objects for user-defined attribute types.
+    """
+    @property
+    def attribute_type(self):
+        """
+        :return: a supported DataJoint attribute type to use; e.g. "longblob", "blob@store"
+        """
+        raise NotImplementedError('Undefined attribute adapter')
+
+    def get(self, value):
+        """
+        convert value retrieved from the the attribute in a table into the adapted type
+        :param value: value from the database
+        :return: object of the adapted type
+        """
+        raise NotImplementedError('Undefined attribute adapter')
+
+    def put(self, obj):
+        """
+        convert an object of the adapted type into a value that DataJoint can store in a table attribute
+        :param object: an object of the adapted type
+        :return: value to store in the database
+        """
+        raise NotImplementedError('Undefined attribute adapter')
+
+
+def get_adapter(context, adapter_name):
+    """
+    Extract the AttributeAdapter object by its name from the context and validate.
+    """
+    if not _support_adapted_types():
+        raise DataJointError('Support for Adapted Attribute types is disabled.')
+    adapter_name = adapter_name.lstrip('<').rstrip('>')
+    try:
+        adapter = context[adapter_name]
+    except KeyError:
+        raise DataJointError(
+            "Attribute adapter '{adapter_name}' is not defined.".format(adapter_name=adapter_name)) from None
+    if not isinstance(adapter, AttributeAdapter):
+        raise DataJointError(
+            "Attribute adapter '{adapter_name}' must be an instance of datajoint.AttributeAdapter".format(
+                adapter_name=adapter_name))
+    if not isinstance(adapter.attribute_type, str) or not re.match(r'^\w', adapter.attribute_type):
+        raise DataJointError("Invalid attribute type {type} in attribute adapter '{adapter_name}'".format(
+            type=adapter.attribute_type, adapter_name=adapter_name))
+    return adapter

datajoint/declare.py

@@ -6,6 +6,7 @@
 import pyparsing as pp
 import logging
 from .errors import DataJointError
+from .attribute_adapter import get_adapter
 
 from .utils import OrderedDict
 
@@ -27,21 +28,24 @@
     INTERNAL_ATTACH=r'attach$',
     EXTERNAL_ATTACH=r'attach@(?P<store>[a-z]\w*)$',
     FILEPATH=r'filepath@(?P<store>[a-z]\w*)$',
-    UUID=r'uuid$').items()}
+    UUID=r'uuid$',
+    ADAPTED=r'<.+>$'
+).items()}
 
-CUSTOM_TYPES = {'UUID', 'INTERNAL_ATTACH', 'EXTERNAL_ATTACH', 'EXTERNAL_BLOB', 'FILEPATH'}  # types stored in attribute comment
+# custom types are stored in attribute comment
+SPECIAL_TYPES = {'UUID', 'INTERNAL_ATTACH', 'EXTERNAL_ATTACH', 'EXTERNAL_BLOB', 'FILEPATH', 'ADAPTED'}
+NATIVE_TYPES = set(TYPE_PATTERN) - SPECIAL_TYPES
 EXTERNAL_TYPES = {'EXTERNAL_ATTACH', 'EXTERNAL_BLOB', 'FILEPATH'}  # data referenced by a UUID in external tables
 SERIALIZED_TYPES = {'EXTERNAL_ATTACH', 'INTERNAL_ATTACH', 'EXTERNAL_BLOB', 'INTERNAL_BLOB'}  # requires packing data
 
-assert set().union(CUSTOM_TYPES, EXTERNAL_TYPES, SERIALIZED_TYPES) <= set(TYPE_PATTERN)
+assert set().union(SPECIAL_TYPES, EXTERNAL_TYPES, SERIALIZED_TYPES) <= set(TYPE_PATTERN)
 
 
-def match_type(datatype):
-    for category, pattern in TYPE_PATTERN.items():
-        match = pattern.match(datatype)
-        if match:
-            return category, match
-    raise DataJointError('Unsupported data types "%s"' % datatype)
+def match_type(attribute_type):
+    try:
+        return next(category for category, pattern in TYPE_PATTERN.items() if pattern.match(attribute_type))
+    except StopIteration:
+        raise DataJointError("Unsupported attribute type {type}".format(type=attribute_type)) from None
 
 
 logger = logging.getLogger(__name__)
@@ -78,7 +82,8 @@ def build_attribute_parser():
     quoted = pp.QuotedString('"') ^ pp.QuotedString("'")
     colon = pp.Literal(':').suppress()
     attribute_name = pp.Word(pp.srange('[a-z]'), pp.srange('[a-z0-9_]')).setResultsName('name')
-    data_type = pp.Combine(pp.Word(pp.alphas) + pp.SkipTo("#", ignore=quoted)).setResultsName('type')
+    data_type = (pp.Combine(pp.Word(pp.alphas) + pp.SkipTo("#", ignore=quoted))
+                 ^ pp.QuotedString('<', endQuoteChar='>', unquoteResults=False)).setResultsName('type')
     default = pp.Literal('=').suppress() + pp.SkipTo(colon, ignore=quoted).setResultsName('default')
     comment = pp.Literal('#').suppress() + pp.restOfLine.setResultsName('comment')
     return attribute_name + pp.Optional(default) + colon + data_type + comment
@@ -168,8 +173,7 @@ def compile_foreign_key(line, context, attributes, primary_key, attr_sql, foreig
             raise DataJointError('Invalid foreign key attributes in "%s"' % line)
         try:
             raise DataJointError('Duplicate attributes "{attr}" in "{line}"'.format(
-                attr=next(attr for attr in result.new_attrs if attr in attributes),
-                line=line))
+                attr=next(attr for attr in result.new_attrs if attr in attributes), line=line))
         except StopIteration:
             pass   # the normal outcome
 
@@ -246,7 +250,7 @@ def prepare_declare(definition, context):
         elif re.match(r'^(unique\s+)?index[^:]*$', line, re.I):   # index
             compile_index(line, index_sql)
         else:
-            name, sql, store = compile_attribute(line, in_key, foreign_key_sql)
+            name, sql, store = compile_attribute(line, in_key, foreign_key_sql, context)
             if store:
                 external_stores.append(store)
             if in_key and name not in primary_key:
@@ -292,10 +296,9 @@ def _make_attribute_alter(new, old, primary_key):
     :param primary_key: primary key attributes
     :return: list of SQL ALTER commands
     """
-
     # parse attribute names
     name_regexp = re.compile(r"^`(?P<name>\w+)`")
-    original_regexp = re.compile(r'COMMENT "\{\s*(?P<name>\w+)\s*\}')
+    original_regexp = re.compile(r'COMMENT "{\s*(?P<name>\w+)\s*}')
     matched = ((name_regexp.match(d), original_regexp.search(d)) for d in new)
     new_names = OrderedDict((d.group('name'), n and n.group('name')) for d, n in matched)
     old_names = [name_regexp.search(d).group('name') for d in old]
@@ -380,13 +383,41 @@ def compile_index(line, index_sql):
         attrs=','.join('`%s`' % a for a in match.attr_list)))
 
 
-def compile_attribute(line, in_key, foreign_key_sql):
+def substitute_special_type(match, category, foreign_key_sql, context):
     """
-    Convert attribute definition from DataJoint format to SQL
+    :param match: dict containing with keys "type" and "comment" -- will be modified in place
+    :param category: attribute type category from TYPE_PATTERN
+    :param foreign_key_sql: list of foreign key declarations to add to
+    :param context: context for looking up user-defined attribute_type adapters
+    """
+    if category == 'UUID':
+        match['type'] = UUID_DATA_TYPE
+    elif category == 'INTERNAL_ATTACH':
+        match['type'] = 'LONGBLOB'
+    elif category in EXTERNAL_TYPES:
+        match['store'] = match['type'].split('@', 1)[1]
+        match['type'] = UUID_DATA_TYPE
+        foreign_key_sql.append(
+            "FOREIGN KEY (`{name}`) REFERENCES `{{database}}`.`{external_table_root}_{store}` (`hash`) "
+            "ON UPDATE RESTRICT ON DELETE RESTRICT".format(external_table_root=EXTERNAL_TABLE_ROOT, **match))
+    elif category == 'ADAPTED':
+        adapter = get_adapter(context, match['type'])
+        match['type'] = adapter.attribute_type
+        category = match_type(match['type'])
+        if category in SPECIAL_TYPES:
+            # recursive redefinition from user-defined datatypes.
+            substitute_special_type(match, category, foreign_key_sql, context)
+    else:
+        assert False, 'Unknown special type'
+
 
+def compile_attribute(line, in_key, foreign_key_sql, context):
+    """
+    Convert attribute definition from DataJoint format to SQL
     :param line: attribution line
     :param in_key: set to True if attribute is in primary key set
-    :param foreign_key_sql:
+    :param foreign_key_sql: the list of foreign key declarations to add to
+    :param context: context in which to look up user-defined attribute type adapterss
     :returns: (name, sql, is_external) -- attribute name and sql code for its declaration
     """
     try:
@@ -412,27 +443,18 @@ def compile_attribute(line, in_key, foreign_key_sql):
             match['default'] = 'NOT NULL'
 
     match['comment'] = match['comment'].replace('"', '\\"')   # escape double quotes in comment
-    category, type_match = match_type(match['type'])
 
     if match['comment'].startswith(':'):
         raise DataJointError('An attribute comment must not start with a colon in comment "{comment}"'.format(**match))
 
-    if category in CUSTOM_TYPES:
+    category = match_type(match['type'])
+    if category in SPECIAL_TYPES:
         match['comment'] = ':{type}:{comment}'.format(**match)  # insert custom type into comment
-        if category == 'UUID':
-            match['type'] = UUID_DATA_TYPE
-        elif category == 'INTERNAL_ATTACH':
-            match['type'] = 'LONGBLOB'
-        elif category in EXTERNAL_TYPES:
-            match['store'] = match['type'].split('@', 1)[1]
-            match['type'] = UUID_DATA_TYPE
-            foreign_key_sql.append(
-                "FOREIGN KEY (`{name}`) REFERENCES `{{database}}`.`{external_table_root}_{store}` (`hash`) "
-                "ON UPDATE RESTRICT ON DELETE RESTRICT".format(external_table_root=EXTERNAL_TABLE_ROOT, **match))
+        substitute_special_type(match, category, foreign_key_sql, context)
 
     if category in SERIALIZED_TYPES and match['default'] not in {'DEFAULT NULL', 'NOT NULL'}:
         raise DataJointError(
-            'The default value for a blob or attachment attributes can only be NULL in:\n%s' % line)
+            'The default value for a blob or attachment attributes can only be NULL in:\n{line}'.format(line=line))
 
     sql = ('`{name}` {type} {default}' + (' COMMENT "{comment}"' if match['comment'] else '')).format(**match)
     return match['name'], sql, match.get('store')

datajoint/fetch.py

@@ -43,35 +43,38 @@ def _get(connection, attr, data, squeeze, download_path):
     """
     if data is None:
         return
+
     extern = connection.schemas[attr.database].external[attr.store] if attr.is_external else None
 
+    adapt = attr.adapter.get if attr.adapter else lambda x: x
+
     if attr.is_filepath:
-        return extern.fget(uuid.UUID(bytes=data))[0]
+        return adapt(extern.fget(uuid.UUID(bytes=data))[0])
 
     if attr.is_attachment:
-        # Steps: 
+        # Steps:
         # 1. peek the filename from the blob without downloading remote
         # 2. check if the file already exists at download_path, verify checksum
         # 3. if exists and checksum passes then return the local filepath
         # 4. Otherwise, download the remote file and return the new filepath
         peek, size = extern.peek(uuid.UUID(bytes=data)) if attr.is_external else (data, len(data))
-        assert size is not None 
+        assert size is not None
         filename = peek.split(b"\0", 1)[0].decode()
         size -= len(filename) + 1
         filepath = os.path.join(download_path, filename)
         if os.path.isfile(filepath) and size == os.path.getsize(filepath):
             local_checksum = hash.uuid_from_file(filepath, filename + '\0')
             remote_checksum = uuid.UUID(bytes=data) if attr.is_external else hash.uuid_from_buffer(data)
-            if local_checksum == remote_checksum: 
-                return filepath  # the existing file is okay
+            if local_checksum == remote_checksum:
+                return adapt(filepath)  # the existing file is okay
         # Download remote attachment
         if attr.is_external:
             data = extern.get(uuid.UUID(bytes=data))
-        return attach.save(data, download_path)  # download file from remote store
+        return adapt(attach.save(data, download_path))  # download file from remote store
 
-    return uuid.UUID(bytes=data) if attr.uuid else (
+    return adapt(uuid.UUID(bytes=data) if attr.uuid else (
             blob.unpack(extern.get(uuid.UUID(bytes=data)) if attr.is_external else data, squeeze=squeeze) 
-            if attr.is_blob else data)
+            if attr.is_blob else data))
 
 
 def _flatten_attribute_list(primary_key, attrs):

datajoint/heading.py

@@ -4,16 +4,17 @@
 import re
 import logging
 from .errors import DataJointError
-from .declare import UUID_DATA_TYPE, CUSTOM_TYPES, TYPE_PATTERN, EXTERNAL_TYPES
+from .declare import UUID_DATA_TYPE, SPECIAL_TYPES, TYPE_PATTERN, EXTERNAL_TYPES, NATIVE_TYPES
 from .utils import OrderedDict
+from .attribute_adapter import get_adapter, AttributeAdapter
 
 
 logger = logging.getLogger(__name__)
 
 default_attribute_properties = dict(    # these default values are set in computed attributes
     name=None, type='expression', in_key=False, nullable=False, default=None, comment='calculated attribute',
     autoincrement=False, numeric=None, string=None, uuid=False, is_blob=False, is_attachment=False, is_filepath=False,
-    is_external=False,
+    is_external=False, adapter=None,
     store=None, unsupported=False, sql_expression=None, database=None, dtype=object)
 
 
@@ -146,7 +147,7 @@ def as_sql(self):
     def __iter__(self):
         return iter(self.attributes)
 
-    def init_from_database(self, conn, database, table_name):
+    def init_from_database(self, conn, database, table_name, context):
         """
         initialize heading from a database table.  The table must exist already.
         """
@@ -211,32 +212,53 @@ def init_from_database(self, conn, database, table_name):
                 numeric=any(TYPE_PATTERN[t].match(attr['type']) for t in ('DECIMAL', 'INTEGER', 'FLOAT')),
                 string=any(TYPE_PATTERN[t].match(attr['type']) for t in ('ENUM', 'TEMPORAL', 'STRING')),
                 is_blob=bool(TYPE_PATTERN['INTERNAL_BLOB'].match(attr['type'])),
-                uuid=False, is_attachment=False, is_filepath=False, store=None, is_external=False, sql_expression=None)
+                uuid=False, is_attachment=False, is_filepath=False, adapter=None,
+                store=None, is_external=False, sql_expression=None)
 
             if any(TYPE_PATTERN[t].match(attr['type']) for t in ('INTEGER', 'FLOAT')):
                 attr['type'] = re.sub(r'\(\d+\)', '', attr['type'], count=1)  # strip size off integers and floats
             attr['unsupported'] = not any((attr['is_blob'], attr['numeric'], attr['numeric']))
             attr.pop('Extra')
 
             # process custom DataJoint types
-            custom_type = re.match(r':(?P<type>[^:]+):(?P<comment>.*)', attr['comment'])
-            if custom_type:
-                attr.update(custom_type.groupdict(), unsupported=False)
+            special = re.match(r':(?P<type>[^:]+):(?P<comment>.*)', attr['comment'])
+            if special:
+                special = special.groupdict()
+                attr.update(special)
+            # process adapted attribute types
+            if special and TYPE_PATTERN['ADAPTED'].match(attr['type']):
+                assert context is not None, 'Declaration context is not set'
+                adapter_name = special['type']
                 try:
-                    category = next(c for c in CUSTOM_TYPES if TYPE_PATTERN[c].match(attr['type']))
+                    attr.update(adapter=get_adapter(context, adapter_name))
+                except DataJointError:
+                    # if no adapter, then delay the error until the first invocation
+                    attr.update(adapter=AttributeAdapter())
+                else:
+                    attr.update(type=attr['adapter'].attribute_type)
+                    if not any(r.match(attr['type']) for r in TYPE_PATTERN.values()):
+                        raise DataJointError(
+                            "Invalid attribute type '{type}' in adapter object <{adapter_name}>.".format(
+                                adapter_name=adapter_name, **attr))
+                    special = not any(TYPE_PATTERN[c].match(attr['type']) for c in NATIVE_TYPES)
+
+            if special:
+                try:
+                    category = next(c for c in SPECIAL_TYPES if TYPE_PATTERN[c].match(attr['type']))
                 except StopIteration:
                     if attr['type'].startswith('external'):
                         raise DataJointError('Legacy datatype `{type}`.'.format(**attr)) from None
                     raise DataJointError('Unknown attribute type `{type}`'.format(**attr)) from None
                 attr.update(
+                    unsupported=False,
                     is_attachment=category in ('INTERNAL_ATTACH', 'EXTERNAL_ATTACH'),
                     is_filepath=category == 'FILEPATH',
                     is_blob=category in ('INTERNAL_BLOB', 'EXTERNAL_BLOB'),  # INTERNAL_BLOB is not a custom type but is included for completeness
                     uuid=category == 'UUID',
                     is_external=category in EXTERNAL_TYPES,
                     store=attr['type'].split('@')[1] if category in EXTERNAL_TYPES else None)
 
-            if attr['in_key'] and (attr['is_blob'] or attr['is_attachment'] or attr['is_filepath']):
+            if attr['in_key'] and any((attr['is_blob'], attr['is_attachment'], attr['is_filepath'])):
                 raise DataJointError('Blob, attachment, or filepath attributes are not allowed in the primary key')
 
             if attr['string'] and attr['default'] is not None and attr['default'] not in sql_literals:
@@ -256,6 +278,11 @@ def init_from_database(self, conn, database, table_name):
                     t = re.sub(r' unsigned$', '', t)   # remove unsigned
                     assert (t, is_unsigned) in numeric_types, 'dtype not found for type %s' % t
                     attr['dtype'] = numeric_types[(t, is_unsigned)]
+
+            if attr['adapter']:
+                # restore adapted type name
+                attr['type'] = adapter_name
+
         self.attributes = OrderedDict(((q['name'], Attribute(**q)) for q in attributes))
 
         # Read and tabulate secondary indexes