Repository: cogent3/cogent3
Pull request #1318 (2023-03-14 22:24:43)

ENH: allow custom handling of the gff attributes field


Patch:
diff --git a/src/cogent3/parse/gff.py b/src/cogent3/parse/gff.py
index d1ad463c9..4a7ef79a1 100644
--- a/src/cogent3/parse/gff.py
+++ b/src/cogent3/parse/gff.py
@@ -1,4 +1,8 @@
-#!/usr/bin/env python
+import collections.abc as abc
+import functools
+import os
+import typing
+
 from pathlib import Path
 
 from cogent3.util.io import open_
@@ -19,29 +23,42 @@ __email__ = "pm67nz@gmail.com"
 __status__ = "Production"
 
 
-def gff_parser(f):
+@functools.singledispatch
+def gff_parser(
+    f: typing.Union[str, os.PathLike, abc.Sequence, typing.IO], attribute_parser=None
+) -> typing.Iterable[dict]:
     """parses a gff file
+
     Parameters
     -----------
     f
         accepts string path or pathlib.Path or file-like object (e.g. StringIO)
+    attribute_parser
+        callback function for custom handling of attributes field. Defaults to
+        a function that converts Attributes field into a dict based on the gff
+        version.
 
     Returns
     -------
     dict
         contains each of the 9 parameters specified by gff3, and comments.
     """
+    yield from _gff_parser(f, attribute_parser=attribute_parser)
 
-    # calling a separate function to ensure file closes correctly
-    f = f if not isinstance(f, Path) else str(f)
-    if isinstance(f, str):
-        with open_(f) as infile:
-            yield from _gff_parser(infile)
-    else:
-        yield from _gff_parser(f)
+
+@gff_parser.register
+def _(f: str, attribute_parser=None) -> typing.Iterable[dict]:
+    with open_(f) as infile:
+        yield from _gff_parser(infile, attribute_parser=attribute_parser)
 
 
-def _gff_parser(f):
+@gff_parser.register
+def _(f: os.PathLike, attribute_parser=None) -> typing.Iterable[dict]:
+    with open_(f) as infile:
+        yield from _gff_parser(infile, attribute_parser=attribute_parser)
+
+
+def _gff_parser(f, attribute_parser=None):
     """parses a gff file"""
 
     gff3_header = "gff-version 3"
@@ -51,10 +68,8 @@ def _gff_parser(f):
         gff3 = gff3_header in f.readline()
         f.seek(0)
 
-    if gff3:
-        attribute_parser = parse_attributes_gff3
-    else:
-        attribute_parser = parse_attributes_gff2
+    if attribute_parser is None:
+        attribute_parser = parse_attributes_gff3 if gff3 else parse_attributes_gff2
 
     for line in f:
         # comments and blank lines
diff --git a/src/cogent3/parse/sequence.py b/src/cogent3/parse/sequence.py
index e8ae305f5..17779afe5 100644
--- a/src/cogent3/parse/sequence.py
+++ b/src/cogent3/parse/sequence.py
@@ -17,7 +17,7 @@ from cogent3.parse import (
     tinyseq,
 )
 from cogent3.parse.record import FileFormatError
-from cogent3.util.io import open_, get_format_suffixes
+from cogent3.util.io import get_format_suffixes, open_
 
 
 __author__ = "Cath Lawrence"


Test patch:
diff --git a/tests/test_parse/test_gff.py b/tests/test_parse/test_gff.py
index 10398a50f..6ac1dc6b7 100644
--- a/tests/test_parse/test_gff.py
+++ b/tests/test_parse/test_gff.py
@@ -7,7 +7,7 @@ from io import StringIO
 from pathlib import Path
 from unittest import TestCase, main
 
-from cogent3.parse.gff import *
+from cogent3.parse.gff import gff_parser, parse_attributes_gff2
 
 
 __author__ = "Matthew Wakefield"
@@ -148,6 +148,12 @@ class GffTest(TestCase):
         # 15 total lines, but 2 comments
         self.assertEqual(i + 1, 15 - 2)
 
+    def test_custom_attr_func(self):
+        """user provided attr parser"""
+        gff3_path = os.path.join("data/c_elegans_WS199_shortened_gff.gff3")
+        for result in gff_parser(gff3_path, attribute_parser=lambda x, y: x):
+            self.assertIsInstance(result["Attributes"], str)
+
 
 if __name__ == "__main__":
     main()