diff --git a/schemainspect/pg/obj.py b/schemainspect/pg/obj.py index bc8639a..bd3b797 100644 --- a/schemainspect/pg/obj.py +++ b/schemainspect/pg/obj.py @@ -1667,11 +1667,14 @@ def filter_schema(self, schema=None, exclude_schema=None): if schema and exclude_schema: raise ValueError("Can only have schema or exclude schema, not both") + if isinstance(exclude_schema, str): + exclude_schema = exclude_schema.split(",") + def equal_to_schema(x): return x.schema == schema def not_equal_to_exclude_schema(x): - return x.schema != exclude_schema + return x.schema not in exclude_schema if schema: comparator = equal_to_schema diff --git a/tests/test_excludeschema.py b/tests/test_excludeschema.py index dca6497..db9c36c 100644 --- a/tests/test_excludeschema.py +++ b/tests/test_excludeschema.py @@ -5,14 +5,14 @@ from .test_all import setup_pg_schema -def asserts_pg_excludedschema(i, schema_names, excludedschema_name): +def asserts_pg_excludedschema(i, schema_names, excludedschemas): schemas = set() for ( prop ) in "schemas relations tables views functions selectables sequences enums constraints".split(): att = getattr(i, prop) for k, v in att.items(): - assert v.schema != excludedschema_name + assert v.schema not in excludedschemas schemas.add(v.schema) assert schemas == set(schema_names) @@ -24,9 +24,17 @@ def test_postgres_inspect_excludeschema(db): s.execute("create schema forthschema;") i = get_inspector(s, exclude_schema="otherschema") asserts_pg_excludedschema( - i, ["public", "forthschema", "thirdschema"], "otherschema" + i, ["public", "forthschema", "thirdschema"], ["otherschema"] ) i = get_inspector(s, exclude_schema="forthschema") asserts_pg_excludedschema( - i, ["public", "otherschema", "thirdschema"], "forthschema" + i, ["public", "otherschema", "thirdschema"], ["forthschema"] + ) + i = get_inspector(s, exclude_schema="forthschema,otherschema") + asserts_pg_excludedschema( + i, ["public", "thirdschema"], ["forthschema", "otherschema"] + ) + i = get_inspector(s, exclude_schema=["forthschema", "otherschema"]) + asserts_pg_excludedschema( + i, ["public", "thirdschema"], ["forthschema", "otherschema"] )