Fix #177: push biotype filter into SQL for genes() / transcripts()

pyensembl/database.py

@@ -504,7 +504,7 @@ def query_one(
 
     @memoize
     def query_feature_values(
-        self, column, feature, distinct=True, contig=None, strand=None
+        self, column, feature, distinct=True, contig=None, strand=None, biotype=None
     ):
         """
         Run a SQL query against the sqlite3 database, filtered
@@ -531,6 +531,16 @@ def query_feature_values(
             query += " AND strand = ?"
             query_params.append(strand)
 
+        if biotype is not None:
+            biotype_column = "%s_biotype" % feature
+            if not self.column_exists(feature, biotype_column):
+                raise ValueError(
+                    "Cannot filter by biotype: table %r has no column %r"
+                    % (feature, biotype_column)
+                )
+            query += " AND %s = ?" % biotype_column
+            query_params.append(biotype)
+
         rows = self.run_sql_query(query, query_params=query_params)
         return [row[0] for row in rows if row is not None]
 

pyensembl/genome.py

@@ -488,7 +488,13 @@ def delete_index_files(self):
             remove(db_path)
 
     def _all_feature_values(
-        self, column, feature, distinct=True, contig=None, strand=None
+        self,
+        column,
+        feature,
+        distinct=True,
+        contig=None,
+        strand=None,
+        biotype=None,
     ):
         """
         Cached lookup of all values for a particular feature property from
@@ -513,6 +519,11 @@ def _all_feature_values(
         strand : str, optional
             Restrict results to "+" or "-" strands
 
+        biotype : str, optional
+            Restrict results to rows whose ``{feature}_biotype`` column
+            matches this value (e.g. ``"protein_coding"``). Only meaningful
+            for the ``gene`` and ``transcript`` features.
+
         Returns a list constructed from query results.
         """
         return self.db.query_feature_values(
@@ -521,6 +532,7 @@ def _all_feature_values(
             distinct=distinct,
             contig=contig,
             strand=strand,
+            biotype=biotype,
         )
 
     def transcript_sequence(self, transcript_id):
@@ -672,10 +684,10 @@ def contigs(self):
     #
     ###################################################
 
-    def genes(self, contig=None, strand=None):
+    def genes(self, contig=None, strand=None, biotype=None):
         """
         Returns all Gene objects in the database. Can be restricted to a
-        particular contig/chromosome and strand by the following arguments:
+        particular contig/chromosome, strand, or biotype.
 
         Parameters
         ----------
@@ -684,8 +696,12 @@ def genes(self, contig=None, strand=None):
 
         strand : str
             Only return genes on this strand.
+
+        biotype : str
+            Only return genes with this Ensembl ``gene_biotype``
+            (e.g. ``"protein_coding"``).
         """
-        gene_ids = self.gene_ids(contig=contig, strand=strand)
+        gene_ids = self.gene_ids(contig=contig, strand=strand, biotype=biotype)
         return [self.gene_by_id(gene_id) for gene_id in gene_ids]
 
     def gene_by_id(self, gene_id):
@@ -818,13 +834,17 @@ def _query_gene_ids(self, property_name, value, feature="gene"):
         )
         return [str(result[0]) for result in results if result[0]]
 
-    def gene_ids(self, contig=None, strand=None):
+    def gene_ids(self, contig=None, strand=None, biotype=None):
         """
-        What are all the gene IDs
-        (optionally restrict to a given chromosome/contig and/or strand)
+        What are all the gene IDs (optionally restrict to a given
+        chromosome/contig, strand, or ``gene_biotype``).
         """
         return self._all_feature_values(
-            column="gene_id", feature="gene", contig=contig, strand=strand
+            column="gene_id",
+            feature="gene",
+            contig=contig,
+            strand=strand,
+            biotype=biotype,
         )
 
     def gene_ids_of_gene_name(self, gene_name):
@@ -860,13 +880,15 @@ def gene_id_of_protein_id(self, protein_id):
     #
     ###################################################
 
-    def transcripts(self, contig=None, strand=None):
+    def transcripts(self, contig=None, strand=None, biotype=None):
         """
         Construct Transcript object for every transcript entry in
-        the database. Optionally restrict to a particular
-        chromosome using the `contig` argument.
+        the database. Optionally restrict to a particular contig, strand,
+        or ``transcript_biotype`` (e.g. ``"protein_coding"``).
         """
-        transcript_ids = self.transcript_ids(contig=contig, strand=strand)
+        transcript_ids = self.transcript_ids(
+            contig=contig, strand=strand, biotype=biotype
+        )
         return [
             self.transcript_by_id(transcript_id) for transcript_id in transcript_ids
         ]
@@ -1002,9 +1024,13 @@ def _query_transcript_ids(self, property_name, value, feature="transcript"):
         )
         return [result[0] for result in results]
 
-    def transcript_ids(self, contig=None, strand=None):
+    def transcript_ids(self, contig=None, strand=None, biotype=None):
         return self._all_feature_values(
-            column="transcript_id", feature="transcript", contig=contig, strand=strand
+            column="transcript_id",
+            feature="transcript",
+            contig=contig,
+            strand=strand,
+            biotype=biotype,
         )
 
     def transcript_ids_of_gene_id(self, gene_id):

pyensembl/version.py

@@ -1,4 +1,4 @@
-__version__ = "2.9.0"
+__version__ = "2.9.1"
 
 def print_version():
     print(f"v{__version__}")

tests/test_biotype_filter.py

@@ -0,0 +1,78 @@
+"""
+Issue #177: biotype filter pushed into the SQL on Genome.genes(),
+Genome.transcripts(), and the underlying *_ids queries.
+"""
+
+import pytest
+
+from .common import eq_
+from .data import (
+    custom_mouse_genome_grcm38_subset,
+    setup_init_custom_mouse_genome,
+)
+
+
+def setup_module(module):
+    setup_init_custom_mouse_genome()
+
+
+def test_transcripts_filter_by_biotype():
+    all_transcripts = custom_mouse_genome_grcm38_subset.transcripts()
+    eq_(len(all_transcripts), 2)
+    biotypes = {t.biotype for t in all_transcripts}
+    # mouse partial fixture has exactly one protein_coding transcript and
+    # one processed_transcript transcript for Cntnap1
+    assert "protein_coding" in biotypes
+    assert biotypes != {"protein_coding"}
+
+    coding = custom_mouse_genome_grcm38_subset.transcripts(
+        biotype="protein_coding"
+    )
+    eq_(len(coding), 1)
+    eq_(coding[0].biotype, "protein_coding")
+
+
+def test_transcript_ids_filter_by_biotype():
+    coding_ids = custom_mouse_genome_grcm38_subset.transcript_ids(
+        biotype="protein_coding"
+    )
+    eq_(coding_ids, ["ENSMUST00000103109"])
+
+
+def test_genes_filter_by_biotype():
+    coding_genes = custom_mouse_genome_grcm38_subset.genes(
+        biotype="protein_coding"
+    )
+    eq_(len(coding_genes), 1)
+    eq_(coding_genes[0].biotype, "protein_coding")
+    # Filtering by a biotype that's absent should return an empty list, not raise.
+    eq_(
+        custom_mouse_genome_grcm38_subset.genes(biotype="miRNA"),
+        [],
+    )
+
+
+def test_biotype_combines_with_contig_and_strand():
+    # protein_coding transcript ENSMUST00000103109 is on chr11, strand +
+    transcripts = custom_mouse_genome_grcm38_subset.transcripts(
+        contig="11", strand="+", biotype="protein_coding"
+    )
+    eq_(len(transcripts), 1)
+    eq_(transcripts[0].id, "ENSMUST00000103109")
+    # wrong strand should yield zero results
+    eq_(
+        custom_mouse_genome_grcm38_subset.transcripts(
+            contig="11", strand="-", biotype="protein_coding"
+        ),
+        [],
+    )
+
+
+def test_biotype_on_feature_without_biotype_column_raises():
+    # The 'exon' table has no exon_biotype column; if a user were to
+    # call _all_feature_values with biotype= for it we should fail loudly
+    # rather than silently return everything.
+    with pytest.raises(ValueError, match="biotype"):
+        custom_mouse_genome_grcm38_subset._all_feature_values(
+            column="exon_id", feature="exon", biotype="protein_coding"
+        )