Fix FunctionMetadata.fields() crash when arg_names is None

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, [])]

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", [])]