user_tables.py

"""
Hosts the table tiers, user tables should be derived from.
"""

import re

from .autopopulate import AutoPopulate
from .errors import DataJointError
from .table import Table
from .utils import ClassProperty, from_camel_case

_base_regexp = r"[a-z][a-z0-9]*(_[a-z][a-z0-9]*)*"

# attributes that trigger instantiation of user classes


supported_class_attrs = {
    "key_source",
    "describe",
    "alter",
    "heading",
    "populate",
    "progress",
    "primary_key",
    "proj",
    "aggr",
    "join",
    "fetch",
    "fetch1",
    "head",
    "tail",
    "descendants",
    "ancestors",
    "parts",
    "parents",
    "children",
    "insert",
    "insert1",
    "update1",
    "drop",
    "drop_quick",
    "delete",
    "delete_quick",
}


class TableMeta(type):
    """
    TableMeta subclasses allow applying some instance methods and properties directly
    at class level. For example, this allows Table.fetch() instead of Table().fetch().
    """

    def __getattribute__(cls, name):
        # trigger instantiation for supported class attrs
        return cls().__getattribute__(name) if name in supported_class_attrs else super().__getattribute__(name)

    def __and__(cls, arg):
        return cls() & arg

    def __xor__(cls, arg):
        return cls() ^ arg

    def __sub__(cls, arg):
        return cls() - arg

    def __neg__(cls):
        return -cls()

    def __mul__(cls, arg):
        return cls() * arg

    def __matmul__(cls, arg):
        return cls() @ arg

    def __add__(cls, arg):
        return cls() + arg

    def __iter__(cls):
        return iter(cls())


class UserTable(Table, metaclass=TableMeta):
    """
    A subclass of UserTable is a dedicated class interfacing a base table.
    UserTable is initialized by the decorator generated by schema().
    """

    # set by @schema
    _connection = None
    _heading = None
    _support = None

    # set by subclass
    tier_regexp = None
    _prefix = None

    @property
    def definition(self):
        """
        :return: a string containing the table definition using the DataJoint DDL.
        """
        raise NotImplementedError('Subclasses of Table must implement the property "definition"')

    @ClassProperty
    def connection(cls):
        return cls._connection

    @ClassProperty
    def table_name(cls):
        """
        :return: the table name of the table formatted for mysql.
        """
        if cls._prefix is None:
            raise AttributeError("Class prefix is not defined!")
        return cls._prefix + from_camel_case(cls.__name__)

    @ClassProperty
    def full_table_name(cls):
        if cls not in {Manual, Imported, Lookup, Computed, Part, UserTable}:
            # for derived classes only
            if cls.database is None:
                raise DataJointError("Class %s is not properly declared (schema decorator not applied?)" % cls.__name__)
            return r"`{0:s}`.`{1:s}`".format(cls.database, cls.table_name)


class Manual(UserTable):
    """
    Inherit from this class if the table's values are entered manually.
    """

    _prefix = r""
    tier_regexp = r"(?P<manual>" + _prefix + _base_regexp + ")"


class Lookup(UserTable):
    """
    Inherit from this class if the table's values are for lookup. This is
    currently equivalent to defining the table as Manual and serves semantic
    purposes only.
    """

    _prefix = "#"
    tier_regexp = r"(?P<lookup>" + _prefix + _base_regexp.replace("TIER", "lookup") + ")"


class Imported(UserTable, AutoPopulate):
    """
    Inherit from this class if the table's values are imported from external data sources.
    The inherited class must at least provide the function `_make_tuples`.
    """

    _prefix = "_"
    tier_regexp = r"(?P<imported>" + _prefix + _base_regexp + ")"

    def drop_quick(self):
        """
        Drop the table and its associated jobs table.
        """
        # Drop the jobs table first if it exists
        if self._jobs_table is not None and self._jobs_table.is_declared:
            self._jobs_table.drop_quick()
        super().drop_quick()


class Computed(UserTable, AutoPopulate):
    """
    Inherit from this class if the table's values are computed from other tables in the schema.
    The inherited class must at least provide the function `_make_tuples`.
    """

    _prefix = "__"
    tier_regexp = r"(?P<computed>" + _prefix + _base_regexp + ")"

    def drop_quick(self):
        """
        Drop the table and its associated jobs table.
        """
        # Drop the jobs table first if it exists
        if self._jobs_table is not None and self._jobs_table.is_declared:
            self._jobs_table.drop_quick()
        super().drop_quick()


class Part(UserTable):
    """
    Inherit from this class if the table's values are details of an entry in another table
    and if this table is populated by the other table. For example, the entries inheriting from
    dj.Part could be single entries of a matrix, while the parent table refers to the entire matrix.
    Part tables are implemented as classes inside classes.
    """

    _connection = None
    _master = None

    tier_regexp = (
        r"(?P<master>"
        + "|".join([c.tier_regexp for c in (Manual, Lookup, Imported, Computed)])
        + r"){1,1}"
        + "__"
        + r"(?P<part>"
        + _base_regexp
        + ")"
    )

    @ClassProperty
    def connection(cls):
        return cls._connection

    @ClassProperty
    def full_table_name(cls):
        return (
            None if cls.database is None or cls.table_name is None else r"`{0:s}`.`{1:s}`".format(cls.database, cls.table_name)
        )

    @ClassProperty
    def master(cls):
        return cls._master

    @ClassProperty
    def table_name(cls):
        return None if cls.master is None else cls.master.table_name + "__" + from_camel_case(cls.__name__)

    def delete(self, force=False, **kwargs):
        """
        Delete contents of a Part table.

        Unless force is True, prohibits direct deletes from parts.

        Args:
            force: If True, allow direct deletion from this Part table.
            **kwargs: Additional keyword arguments passed to Table.delete(),
                such as transaction, safemode, and force_masters.

        Raises:
            DataJointError: If force is False (default).
        """
        if force:
            super().delete(force_parts=True, **kwargs)
        else:
            raise DataJointError("Cannot delete from a Part directly. Delete from master instead")

    def drop(self, force=False):
        """
        unless force is True, prohibits direct deletes from parts.
        """
        if force:
            super().drop()
        else:
            raise DataJointError("Cannot drop a Part directly.  Delete from master instead")

    def alter(self, prompt=True, context=None):
        # without context, use declaration context which maps master keyword to master table
        super().alter(prompt=prompt, context=context or self.declaration_context)


user_table_classes = (Manual, Lookup, Computed, Imported, Part)


class _AliasNode:
    """
    special class to indicate aliased foreign keys
    """

    pass


def _get_tier(table_name):
    """given the table name, return the user table class."""
    if not table_name.startswith("`"):
        return _AliasNode
    else:
        try:
            return next(tier for tier in user_table_classes if re.fullmatch(tier.tier_regexp, table_name.split("`")[-2]))
        except StopIteration:
            return None