HStoreRef: a simple annotation expression to select individual hstore keys

Author: Photonios

psqlextra/expressions.py

@@ -54,3 +54,35 @@ def relabeled_clone(self, relabels):
             self.hstore_key,
             self.output_field
         )
+
+
+class HStoreRef(expressions.F):
+    """Inline reference to a HStore key.
+
+    Allows selecting individual keys in annotations.
+    """
+
+    def __init__(self, name: str, key: str):
+        """Initializes a new instance of :see:HStoreRef.
+
+        Arguments:
+            name:
+                The name of the column/field to resolve.
+
+            key:
+                The name of the HStore key to select.
+        """
+
+        super().__init__(name)
+        self.key = key
+
+    def resolve_expression(self, *args, **kwargs) -> HStoreColumn:
+        """Resolves the expression into a :see:HStoreColumn expression."""
+
+        original_expression = super().resolve_expression(*args, **kwargs)
+        expression = HStoreColumn(
+            original_expression.alias,
+            original_expression.target,
+            self.key
+        )
+        return expression

tests/test_hstore_field.py

@@ -1,6 +1,7 @@
 from django.db import models
 
 from psqlextra import HStoreField
+from psqlextra.expressions import HStoreRef
 
 from .fake_model import get_fake_model
 
@@ -46,3 +47,31 @@ def test_values():
     result = list(model.objects.values_list('title__en', 'title__ar'))[0]
     assert result[0] == obj.title['en']
     assert result[1] == obj.title['ar']
+
+    result = list(model.objects.values_list('title__en', 'title__ar'))[0]
+
+def test_annotate_ref():
+    """Tests whether annotating using a :see:HStoreRef expression
+    works correctly.
+
+    This allows you to select an individual hstore key."""
+
+    model_fk = get_fake_model({
+        'title': HStoreField(),
+    })
+
+    model = get_fake_model({
+        'fk': models.ForeignKey(model_fk)
+    })
+
+    fk = model_fk.objects.create(title={'en': 'english', 'ar': 'arabic'})
+    obj = model.objects.create(fk=fk)
+
+    queryset = (
+        model.objects
+        .annotate(english_title=HStoreRef('fk__title', 'en'))
+        .values('english_title')
+        .first()
+    )
+
+    assert queryset['english_title'] == 'english'