feat: add Pangolineage_query

GenSpectrum · Apr 27, 2023 · c159ace · c159ace
1 parent 835b9c1
commit c159ace
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 5 deletions.
diff --git a/lapis2/src/main/antlr/org/genspectrum/lapis/model/variantqueryparser/VariantQuery.g4 b/lapis2/src/main/antlr/org/genspectrum/lapis/model/variantqueryparser/VariantQuery.g4
@@ -14,14 +14,21 @@ expr:
 
 single:
   nucleotide_mutation
+  | pangolineage_query
   ;
 
 nucleotide_mutation : nucleotide_symbol? position ambigous_nucleotide_symbol?;
-
 position: NUMBER+;
 nucleotide_symbol: A | C | G | T;
 ambigous_nucleotide_symbol: nucleotide_symbol | M | R | W | S | Y | K | V | H | D | B | N | MINUS | DOT;
 
+pangolineage_query: pangolineage pangolineage_include_sublineages?;
+pangolineage: pangolineage_character pangolineage_character? pangolineage_character? pangolineage_number_component*;
+pangolineage_character: A | B | C | D | E | F | G | H | I | J | K | L | M | N | O | P | Q | R | S | T | U | V | W | X | Y | Z;
+pangolineage_number_component: '.' NUMBER NUMBER? NUMBER?;
+pangolineage_include_sublineages: DOT? ASTERISK;
+
+
 // lexer rules
 
 A: 'A';
@@ -52,6 +59,7 @@ Y: 'Y';
 Z: 'Z';
 MINUS: '-';
 DOT: '.';
+ASTERISK: '*';
 
 NUMBER: [0-9];
 WHITESPACE: [ \r\n\t] -> skip;
diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/model/SiloFilterExpressionMapper.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/model/SiloFilterExpressionMapper.kt
@@ -18,7 +18,10 @@ data class SequenceFilterValue(val type: SequenceFilterFieldType, val value: Str
 typealias SequenceFilterFieldName = String
 
 @Component
-class SiloFilterExpressionMapper(private val allowedSequenceFilterFields: SequenceFilterFields) {
+class SiloFilterExpressionMapper(
+    private val allowedSequenceFilterFields: SequenceFilterFields,
+    private val variantQueryFacade: VariantQueryFacade,
+) {
     fun map(sequenceFilters: Map<String, String>): SiloFilterExpression {
         if (sequenceFilters.isEmpty()) {
             return True
@@ -82,7 +85,7 @@ class SiloFilterExpressionMapper(private val allowedSequenceFilterFields: Sequen
             )
         }
 
-        return VariantQueryFacade().map(variantQuery)
+        return variantQueryFacade.map(variantQuery)
     }
 
     private fun mapToDateBetweenFilter(

diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/model/VariantQueryCustomListener.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/model/VariantQueryCustomListener.kt
@@ -6,12 +6,14 @@ import VariantQueryParser.MaybeContext
 import VariantQueryParser.NotContext
 import VariantQueryParser.Nucleotide_mutationContext
 import VariantQueryParser.OrContext
+import VariantQueryParser.Pangolineage_queryContext
 import org.antlr.v4.runtime.tree.ParseTreeListener
 import org.genspectrum.lapis.silo.And
 import org.genspectrum.lapis.silo.Maybe
 import org.genspectrum.lapis.silo.Not
 import org.genspectrum.lapis.silo.NucleotideSymbolEquals
 import org.genspectrum.lapis.silo.Or
+import org.genspectrum.lapis.silo.PangoLineageEquals
 import org.genspectrum.lapis.silo.SiloFilterExpression
 
 class VariantQueryCustomListener : VariantQueryBaseListener(), ParseTreeListener {
@@ -26,12 +28,23 @@ class VariantQueryCustomListener : VariantQueryBaseListener(), ParseTreeListener
             return
         }
         val position = ctx.position().text.toInt()
-        val secondSymbol = if (ctx.ambigous_nucleotide_symbol() != null) ctx.ambigous_nucleotide_symbol().text else "-"
+        val secondSymbol = ctx.ambigous_nucleotide_symbol()?.text ?: "-"
 
         val expr = NucleotideSymbolEquals(position, secondSymbol)
         expressionStack.addLast(expr)
     }
 
+    override fun enterPangolineage_query(ctx: Pangolineage_queryContext?) {
+        if (ctx == null) {
+            return
+        }
+        val pangolineage = ctx.pangolineage().text
+        val includeSublineages = ctx.pangolineage_include_sublineages() != null
+
+        val expr = PangoLineageEquals(pangolineage, includeSublineages)
+        expressionStack.addLast(expr)
+    }
+
     override fun exitAnd(ctx: AndContext?) {
         val children = listOf(expressionStack.removeLast(), expressionStack.removeLast()).reversed()
         expressionStack.addLast(And(children))

diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/model/VariantQueryFacade.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/model/VariantQueryFacade.kt
@@ -6,7 +6,9 @@ import org.antlr.v4.runtime.CharStreams
 import org.antlr.v4.runtime.CommonTokenStream
 import org.antlr.v4.runtime.tree.ParseTreeWalker
 import org.genspectrum.lapis.silo.SiloFilterExpression
+import org.springframework.stereotype.Component
 
+@Component
 class VariantQueryFacade {
 
     fun map(variantQuery: String): SiloFilterExpression {

diff --git a/lapis2/src/test/kotlin/org/genspectrum/lapis/model/SiloFilterExpressionMapperTest.kt b/lapis2/src/test/kotlin/org/genspectrum/lapis/model/SiloFilterExpressionMapperTest.kt
@@ -37,7 +37,7 @@ class SiloFilterExpressionMapperTest {
 
     @BeforeEach
     fun setup() {
-        underTest = SiloFilterExpressionMapper(sequenceFilterFields)
+        underTest = SiloFilterExpressionMapper(sequenceFilterFields, VariantQueryFacade())
     }
 
     @Test

diff --git a/lapis2/src/test/kotlin/org/genspectrum/lapis/model/VariantQueryFacadeTest.kt b/lapis2/src/test/kotlin/org/genspectrum/lapis/model/VariantQueryFacadeTest.kt
@@ -5,6 +5,7 @@ import org.genspectrum.lapis.silo.Maybe
 import org.genspectrum.lapis.silo.Not
 import org.genspectrum.lapis.silo.NucleotideSymbolEquals
 import org.genspectrum.lapis.silo.Or
+import org.genspectrum.lapis.silo.PangoLineageEquals
 import org.hamcrest.MatcherAssert
 import org.hamcrest.Matchers
 import org.junit.jupiter.api.BeforeEach
@@ -117,4 +118,24 @@ class VariantQueryFacadeTest {
         val expectedResult = Maybe(NucleotideSymbolEquals(300, "G"))
         MatcherAssert.assertThat(result, Matchers.equalTo(expectedResult))
     }
+
+    @Test
+    fun `given a variant variantQuery with a 'Pangolineage' expression the map should return the corresponding SiloQuery`() {
+        val variantQuery = "A.1.2.3"
+
+        val result = underTest.map(variantQuery)
+
+        val expectedResult = PangoLineageEquals("A.1.2.3", false)
+        MatcherAssert.assertThat(result, Matchers.equalTo(expectedResult))
+    }
+
+    @Test
+    fun `given a variant variantQuery with a 'Pangolineage' expression including sublineages the map should return the corresponding SiloQuery`() {
+        val variantQuery = "A.1.2.3*"
+
+        val result = underTest.map(variantQuery)
+
+        val expectedResult = PangoLineageEquals("A.1.2.3", true)
+        MatcherAssert.assertThat(result, Matchers.equalTo(expectedResult))
+    }
 }