codekingpro/portable-devtools

1from collections import namedtuple2 3_ColumnMetadata = namedtuple(4    "ColumnMetadata", ["name", "datatype", "foreignkeys", "default",5                       "has_default"]6)7 8 9def ColumnMetadata(name, datatype, foreignkeys=None, default=None,10                   has_default=False):11    return _ColumnMetadata(name, datatype, foreignkeys or [], default,12                           has_default)13 14 15ForeignKey = namedtuple(16    "ForeignKey",17    [18        "parentschema",19        "parenttable",20        "parentcolumn",21        "childschema",22        "childtable",23        "childcolumn",24    ],25)26TableMetadata = namedtuple("TableMetadata", "name columns")27 28 29def parse_defaults(defaults_string):30    """Yields default values for a function, given the string provided by31    pg_get_expr(pg_catalog.pg_proc.proargdefaults, 0)"""32    if not defaults_string:33        return34    current = ""35    in_quote = None36    for char in defaults_string:37        if current == "" and char == " ":38            # Skip space after comma separating default expressions39            continue40        if char == '"' or char == "'":41            if in_quote and char == in_quote:42                # End quote43                in_quote = None44            elif not in_quote:45                # Begin quote46                in_quote = char47        elif char == "," and not in_quote:48            # End of expression49            yield current50            current = ""51            continue52        current += char53    yield current54 55 56class FunctionMetadata:57    def __init__(58        self,59        schema_name,60        func_name,61        arg_names,62        arg_types,63        arg_modes,64        return_type,65        is_aggregate,66        is_window,67        is_set_returning,68        is_extension,69        arg_defaults,70    ):71        """Class for describing a postgresql function"""72 73        self.schema_name = schema_name74        self.func_name = func_name75 76        self.arg_modes = tuple(arg_modes) if arg_modes else None77        self.arg_names = tuple(arg_names) if arg_names else None78 79        # Be flexible in not requiring arg_types -- use None as a placeholder80        # for each arg. (Used for compatibility with old versions of postgresql81        # where such info is hard to get.82        if arg_types:83            self.arg_types = tuple(arg_types)84        elif arg_modes:85            self.arg_types = tuple([None] * len(arg_modes))86        elif arg_names:87            self.arg_types = tuple([None] * len(arg_names))88        else:89            self.arg_types = None90 91        self.arg_defaults = tuple(parse_defaults(arg_defaults))92 93        self.return_type = return_type.strip()94        self.is_aggregate = is_aggregate95        self.is_window = is_window96        self.is_set_returning = is_set_returning97        self.is_extension = bool(is_extension)98        self.is_public = self.schema_name and self.schema_name == "public"99 100    def __eq__(self, other):101        return isinstance(other, self.__class__) and \102            self.__dict__ == other.__dict__103 104    def __ne__(self, other):105        return not self.__eq__(other)106 107    def _signature(self):108        return (109            self.schema_name,110            self.func_name,111            self.arg_names,112            self.arg_types,113            self.arg_modes,114            self.return_type,115            self.is_aggregate,116            self.is_window,117            self.is_set_returning,118            self.is_extension,119            self.arg_defaults,120        )121 122    def __hash__(self):123        return hash(self._signature())124 125    def __repr__(self):126        return (127            "%s(schema_name=%r, func_name=%r, arg_names=%r, "128            "arg_types=%r, arg_modes=%r, return_type=%r, is_aggregate=%r, "129            "is_window=%r, is_set_returning=%r, is_extension=%r, "130            "arg_defaults=%r)"131        ) % ((self.__class__.__name__,) + self._signature())132 133    def has_variadic(self):134        return self.arg_modes and \135            any(arg_mode == "v" for arg_mode in self.arg_modes)136 137    def args(self):138        """Returns a list of input-parameter ColumnMetadata namedtuples."""139        if not self.arg_names:140            return []141        modes = self.arg_modes or ["i"] * len(self.arg_names)142        args = [143            (name, typ)144            for name, typ, mode in zip(self.arg_names, self.arg_types, modes)145            if mode in ("i", "b", "v")  # IN, INOUT, VARIADIC146        ]147 148        def arg(name, typ, num):149            num_args = len(args)150            num_defaults = len(self.arg_defaults)151            has_default = num + num_defaults >= num_args152            default = (153                self.arg_defaults[num - num_args + num_defaults]154                if has_default155                else None156            )157            return ColumnMetadata(name, typ, [], default, has_default)158 159        return [arg(name, typ, num) for num, (name, typ) in enumerate(args)]160 161    def fields(self):162        """Returns a list of output-field ColumnMetadata namedtuples"""163 164        if self.return_type.lower() == "void":165            return []166        elif not self.arg_modes:167            # For functions  without output parameters, the function name168            # is used as the name of the output column.169            # E.g. 'SELECT unnest FROM unnest(...);'170            return [ColumnMetadata(self.func_name, self.return_type, [])]171 172        return [173            ColumnMetadata(name, typ, [])174            for name, typ, mode in zip(self.arg_names, self.arg_types,175                                       self.arg_modes)176            if mode in ("o", "b", "t")177        ]  # OUT, INOUT, TABLE178