diff --git a/phykit/services/alignment/faidx.py b/phykit/services/alignment/faidx.py index 03f9d57..1bf7f9a 100644 --- a/phykit/services/alignment/faidx.py +++ b/phykit/services/alignment/faidx.py @@ -1,3 +1,5 @@ +from typing import Dict + from Bio import SeqIO from .base import Alignment @@ -7,11 +9,13 @@ class Faidx(Alignment): def __init__(self, args) -> None: super().__init__(**self.process_args(args)) - def run(self): + def run(self) -> None: record_dict = SeqIO.index(self.fasta, "fasta") - entry = self.entry.split(',') - for e in entry: - print(f">{record_dict[e].name}\n{record_dict[e].seq}") - def process_args(self, args): + # Split entries and iterate + for e in map(str.strip, self.entry.split(',')): + record = record_dict[e] + print(f">{record.name}\n{record.seq}") + + def process_args(self, args) -> Dict[str, str]: return dict(fasta=args.fasta, entry=args.entry) diff --git a/tests/integration/alignment/test_faidx.py b/tests/integration/alignment/test_faidx.py index b5c34c8..987018c 100644 --- a/tests/integration/alignment/test_faidx.py +++ b/tests/integration/alignment/test_faidx.py @@ -2,7 +2,6 @@ import sys from mock import patch, call from pathlib import Path -from textwrap import dedent from phykit.phykit import Phykit @@ -10,7 +9,7 @@ @pytest.mark.integration -class TestGCContent(object): +class TestFaidx(object): @patch("builtins.print") def test_faidx(self, mocked_print): expected_result = ">1\nA-GTAT"