implement feature type figure

oschwengers · Nov 7, 2024 · ba22dac · ba22dac
1 parent 160fcbb
commit ba22dac
Show file tree

Hide file tree

Showing 3 changed files with 200 additions and 113 deletions.
diff --git a/bakta/io/insdc.py b/bakta/io/insdc.py
@@ -18,9 +18,7 @@
 log = logging.getLogger('INSDC')
 
 
-def write_features(data: dict, features: Sequence[dict], genbank_output_path: Path, embl_output_path: Path):
-    log.debug('prepare: genbank=%s, embl=%s', genbank_output_path, embl_output_path)
-
+def build_biopython_sequence_list(data: dict, features: Sequence[dict]):
     sequence_list = []
     for seq in data['sequences']:
         sequence_features = [feat for feat in features if feat['sequence'] == seq['id']]
@@ -275,16 +273,20 @@ def write_features(data: dict, features: Sequence[dict], genbank_output_path: Pa
                 seq_feature_list.append(acc_feature)
         sequence_record.features = seq_feature_list
         sequence_list.append(sequence_record)
+    return sequence_list
+
 
+def write_features(data: dict, features: Sequence[dict], genbank_output_path: Path, embl_output_path: Path):
+    log.debug('prepare: genbank=%s, embl=%s', genbank_output_path, embl_output_path)
+
+    sequence_list = build_biopython_sequence_list(data, features)
     with genbank_output_path.open('wt', encoding='utf-8') as fh:
         log.info('write GenBank: path=%s', genbank_output_path)
         SeqIO.write(sequence_list, fh, format='genbank')
 
     with embl_output_path.open('wt', encoding='utf-8') as fh:
         log.info('write EMBL: path=%s', embl_output_path)
         SeqIO.write(sequence_list, fh, format='embl')
-
-    return sequence_list
 
 
 def select_ncrna_class(feature: dict) -> str:

diff --git a/bakta/main.py b/bakta/main.py
@@ -553,7 +553,7 @@ def main():
     print('\tINSDC GenBank & EMBL...')
     genbank_path = cfg.output_path.joinpath(f'{cfg.prefix}.gbff')
     embl_path = cfg.output_path.joinpath(f'{cfg.prefix}.embl')
-    biopython_sequences = insdc.write_features(data, features, genbank_path, embl_path)
+    insdc.write_features(data, features, genbank_path, embl_path)
 
     print('\tgenome sequences...')
     fna_path = cfg.output_path.joinpath(f'{cfg.prefix}.fna')
@@ -575,7 +575,7 @@ def main():
         print('\tskip generation of circular genome plot...')
     else:
         print('\tcircular genome plot...')
-        plot.write(biopython_sequences, sequences, cfg.output_path)
+        plot.write(data, features, cfg.output_path)
 
     if(cfg.skip_cds is False):
         hypotheticals = [feat for feat in features if feat['type'] == bc.FEATURE_CDS and 'hypothetical' in feat]