intersection memberships

entity/constants.py

@@ -0,0 +1,4 @@
+
+class MembershipType:
+    UNION = 'UNION'
+    INTERSECTION = 'INTERSECTION'

entity/migrations/0002_entitygroup_membership_type.py

@@ -0,0 +1,18 @@
+# Generated by Django 4.2.4 on 2023-08-29 13:41
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('entity', '0001_0010_squashed'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='entitygroup',
+            name='membership_type',
+            field=models.CharField(choices=[('UNION', 'Union'), ('INTERSECTION', 'Intersection')], default='UNION'),
+        ),
+    ]

entity/models.py

@@ -9,6 +9,8 @@
 from python3_utils import compare_on_attr
 from functools import reduce
 
+from entity.constants import MembershipType
+
 
 class AllEntityKindManager(ActivatableManager):
     """
@@ -363,6 +365,13 @@ class EntityGroup(models.Model):
 
     objects = EntityGroupManager()
 
+    membership_type_choices = [
+        (MembershipType.UNION, 'Union'),
+        (MembershipType.INTERSECTION, 'Intersection'),
+    ]
+
+    membership_type = models.CharField(choices=membership_type_choices, default=MembershipType.UNION)
+
     def all_entities(self, is_active=True):
         """
         Return all the entities in the group.
@@ -405,16 +414,26 @@ def get_all_entities(self, membership_cache=None, entities_by_kind=None, return_
 
             # Loop over each membership in this group
             for entity_id, entity_kind_id in membership_cache[self.id]:
+                entity_ids_to_apply = set()
                 if entity_id:
                     if entity_kind_id:
                         # All sub entities of this kind under this entity
-                        entity_ids.update(entities_by_kind[entity_kind_id][entity_id])
+                        entity_ids_to_apply.update(entities_by_kind[entity_kind_id][entity_id])
                     else:
                         # Individual entity
-                        entity_ids.add(entity_id)
+                        entity_ids_to_apply.add(entity_id)
                 else:
                     # All entities of this kind
-                    entity_ids.update(entities_by_kind[entity_kind_id]['all'])
+                    entity_ids_to_apply.update(entities_by_kind[entity_kind_id]['all'])
+
+                # Check membership type
+                if self.membership_type == MembershipType.UNION:
+                    entity_ids.update(entity_ids_to_apply)
+                elif self.membership_type == MembershipType.INTERSECTION:
+                    if not entity_ids:
+                        entity_ids.update(entity_ids_to_apply)
+                    else:
+                        entity_ids = entity_ids.intersection(entity_ids_to_apply)
 
         # Check if a queryset needs to be returned
         if return_models:

entity/tests/model_tests.py

@@ -5,10 +5,10 @@
 
 from entity.signal_handlers import turn_off_syncing, turn_on_syncing
 
+from entity.constants import MembershipType
 from entity.models import (
     Entity, EntityKind, EntityRelationship, EntityGroup, EntityGroupMembership, get_entities_by_kind
 )
-
 from entity.tests.models import Account, Team, TeamGroup, Competitor
 from entity.tests.utils import EntityTestCase
 
@@ -758,6 +758,55 @@ def setUp(self):
 
         self.group = G(EntityGroup)
 
+    def test_membership_type_intersection(self):
+        """
+        Given two memberships of entities under different entity kinds, verify that only the intersection is returned
+        instead of the union.
+
+        This test sets up:
+        - 5 sub entities under super 1
+        - 5 sub entities under super 2
+        - 3 sub entities under both
+        """
+        super_entity_kind1 = G(EntityKind)
+        super_entity_kind2 = G(EntityKind)
+        sub_entity_kind = G(EntityKind)
+        super_entity1 = G(Entity, entity_kind=super_entity_kind1)
+        super_entity2 = G(Entity, entity_kind=super_entity_kind2)
+        sub_entities1 = [
+            G(Entity, entity_kind=sub_entity_kind)
+            for _ in range(5)
+        ]
+        sub_entities2 = [
+            G(Entity, entity_kind=sub_entity_kind)
+            for _ in range(5)
+        ]
+
+        # Create the relationships
+        for entity in sub_entities1:
+            G(EntityRelationship, sub_entity=entity, super_entity=super_entity1)
+        for entity in sub_entities2:
+            G(EntityRelationship, sub_entity=entity, super_entity=super_entity2)
+
+        # Create the intersection relationships
+        G(EntityRelationship, sub_entity=sub_entities1[0], super_entity=super_entity2)
+        G(EntityRelationship, sub_entity=sub_entities1[1], super_entity=super_entity2)
+        G(EntityRelationship, sub_entity=sub_entities1[2], super_entity=super_entity2)
+
+        # Create the entity group
+        entity_group = G(EntityGroup, membership_type=MembershipType.INTERSECTION)
+
+        # Create the memberships -- two memberships of all subs under a kind
+        G(EntityGroupMembership, entity_group=entity_group, sub_entity_kind=sub_entity_kind, entity=super_entity1)
+        G(EntityGroupMembership, entity_group=entity_group, sub_entity_kind=sub_entity_kind, entity=super_entity2)
+
+        entity_ids = entity_group.get_all_entities()
+        self.assertEqual(entity_ids, set([
+            sub_entities1[0].id,
+            sub_entities1[1].id,
+            sub_entities1[2].id,
+        ]))
+
     def test_individual_entities_returned(self):
         e = self.super_entities[0]
         G(EntityGroupMembership, entity_group=self.group, entity=e, sub_entity_kind=None)

entity/version.py

@@ -1 +1 @@
-__version__ = '6.1.1'
+__version__ = '6.2.0'

release_notes.md

@@ -1,5 +1,7 @@
 ## Release Notes
 
+- 6.2.0:
+    - Add support for intersection type memberships
 - 6.1.1:
     - django support for 4.2
     - drop django 2.2