diff --git a/pgcli/packages/parseutils/meta.py b/pgcli/packages/parseutils/meta.py index df41cf4ee..73decbff4 100644 --- a/pgcli/packages/parseutils/meta.py +++ b/pgcli/packages/parseutils/meta.py @@ -157,8 +157,14 @@ def fields(self): # E.g. 'SELECT unnest FROM unnest(...);' return [ColumnMetadata(self.func_name, self.return_type, [])] - return [ + # arg_modes being truthy doesn't guarantee arg_names/arg_types are + # populated too (e.g. an unnamed variadic parameter or a TABLE(...) + # return without argument names). + fields = [ ColumnMetadata(name, typ, []) - for name, typ, mode in zip(self.arg_names, self.arg_types, self.arg_modes) - if mode in ("o", "b", "t") - ] # OUT, INOUT, TABLE + for name, typ, mode in zip(self.arg_names or [], self.arg_types or [], self.arg_modes) + if mode in ("o", "b", "t") # OUT, INOUT, TABLE + ] + # Without any usable output column, fall back to the function name, + # as for functions declared without output parameters. + return fields or [ColumnMetadata(self.func_name, self.return_type, [])] diff --git a/tests/parseutils/test_function_metadata.py b/tests/parseutils/test_function_metadata.py index c4000ab1c..938e2a5ea 100644 --- a/tests/parseutils/test_function_metadata.py +++ b/tests/parseutils/test_function_metadata.py @@ -1,4 +1,5 @@ -from pgcli.packages.parseutils.meta import FunctionMetadata +from pgcli.packages.parseutils.meta import ColumnMetadata, FunctionMetadata +from pgcli.pgcompleter import generate_alias def test_function_metadata_eq(): @@ -11,3 +12,25 @@ def test_function_metadata_eq(): assert not (f1 == f3) assert hash(f1) == hash(f2) assert hash(f1) != hash(f3) + + +def test_function_metadata_fields_with_variadic_and_no_arg_names(): + # Regression test: arg_modes being truthy doesn't guarantee arg_names is + # populated (e.g. an unnamed variadic parameter). fields() used to crash + # with "'NoneType' object is not iterable". + f = FunctionMetadata("s", "labels", None, ["text[]"], ["v"], "hstore", False, False, False, False, None) + assert f.fields() == [ColumnMetadata("labels", "hstore", [])] + + +def test_function_metadata_fields_table_mode_with_no_arg_names(): + # Without argument names there is no output column name to offer, so the + # function name is used and generate_alias() gets a real string. + f = FunctionMetadata("s", "f", None, ["int4", "text"], ["t", "t"], "record", False, False, True, False, None) + fields = f.fields() + assert fields == [ColumnMetadata("f", "record", [])] + assert all(generate_alias(field.name) for field in fields) + + +def test_function_metadata_fields_table_mode_with_arg_names(): + f = FunctionMetadata("s", "f", ["a", "b"], ["int4", "text"], ["t", "t"], "record", False, False, True, False, None) + assert f.fields() == [ColumnMetadata("a", "int4", []), ColumnMetadata("b", "text", [])]