Enable entity category import from module search path

This enhancement causes an entity category import to first be tried from
the general module search path, and if that fails then to fall back to
the current default of importing saml2.entity_category.<module>. This
allows deployers to overlay their own customized versions of entity
category modules like edugain.py that contain CoCo. This is helpful
since the list of attributes to be included as part of the entity
category may not be globally the same for all deployments. Such is the
case with CoCo where the list of attributes changes from federation to
federation and deployment to deployment.

1 files changed, 5 insertions, 2 deletions

diff --git a/src/saml2/assertion.py b/src/saml2/assertion.py
index 8984db59..ed043d9a 100644
--- a/src/saml2/assertion.py
+++ b/src/saml2/assertion.py
@@ -353,8 +353,11 @@ class Policy(object):
             else:
                 ecs = []
                 for cat in items:
-                    _mod = importlib.import_module(
-                        "saml2.entity_category.%s" % cat)
+                    try:
+                        _mod = importlib.import_module(cat)
+                    except ImportError:
+                        _mod = importlib.import_module(
+                            "saml2.entity_category.%s" % cat)
                     _ec = {}
                     for key, items in _mod.RELEASE.items():
                         alist = [k.lower() for k in items]