Coverage for rest_api/data_entry/gbk_import.py: 18%
247 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-10 03:22 +0000
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-10 03:22 +0000
1import calendar
2from datetime import datetime
3import os
4import pathlib
6from Bio import SeqFeature
7from Bio import SeqIO
8from Bio import SeqRecord
9from django.core.files.uploadedfile import InMemoryUploadedFile
10from django.db import transaction
12from rest_api.models import CDS
13from rest_api.models import CDSSegment
14from rest_api.models import Gene
15from rest_api.models import GeneSegment
16from rest_api.models import Peptide
17from rest_api.models import PeptideSegment
18from rest_api.models import Reference
19from rest_api.models import Replicon
20from rest_api.serializers import CDSSegmentSerializer
21from rest_api.serializers import CDSSerializer
22from rest_api.serializers import find_or_create
23from rest_api.serializers import GeneSegmentSerializer
24from rest_api.serializers import GeneSerializer
25from rest_api.serializers import PeptideSegmentSerializer
26from rest_api.serializers import PeptideSerializer
27from rest_api.serializers import ReferenceSerializer
28from rest_api.serializers import RepliconSerializer
29from sonar_backend.settings import SONAR_DATA_ENTRY_FOLDER
32def import_gbk_file(uploaded_files: list[InMemoryUploadedFile], translation_id: int):
33 """
34 Import GenBank file. or Import multiple GenBank files as segments.
36 Args:
37 - uploaded_files (list): List of InMemoryUploadedFile objects.
38 - translation_id (int): ID for translation.
39 """
40 records = []
41 file_paths = []
42 for uploaded_file in uploaded_files:
43 file_path = _temp_save_file(uploaded_file)
44 file_paths.append(str(file_path))
45 records.extend(list(SeqIO.parse(file_path, "genbank")))
47 file_path_str = ", ".join(
48 file_paths
49 ) # join multiple file paths into a comma-space separated string
50 records: list[SeqRecord.SeqRecord]
51 reference = _put_reference_from_record(records[0], translation_id, file_path_str)
52 with transaction.atomic():
53 # record = complete gbk file
54 for record in records:
55 source_features = list(
56 filter(lambda x: x.type == "source", record.features)
57 )
58 if len(source_features) != 1:
59 raise ValueError("Expecting exactly one source feature.")
60 source_feature = source_features[0]
61 source_feature: SeqFeature.SeqFeature
62 if source_feature.location is None:
63 raise ValueError("No location information found for source feature.")
64 replicon_data = {
65 "accession": f"{record.name}.{record.annotations['sequence_version']}",
66 "description": record.description,
67 "length": int(source_feature.location.end)
68 - int(source_feature.location.start),
69 "sequence": str(source_feature.extract(record.seq)),
70 "reference": reference.id,
71 }
72 # If the record has a segment number at the top level, use that
73 # else look at the source feature for segment number
74 if "segment_number" in record.annotations:
75 replicon_data["segment_number"] = record.annotations["segment_number"]
76 else:
77 source_features = [f for f in record.features if f.type == "source"]
79 if (
80 len(source_features) == 1
81 and "segment" in source_features[0].qualifiers
82 ):
83 replicon_data["segment_number"] = source_features[0].qualifiers[
84 "segment"
85 ][0]
87 replicon = find_or_create(replicon_data, Replicon, RepliconSerializer)
88 # features with gene qualifier
89 gene_features = [
90 f
91 for f in record.features
92 if f.type == "gene" and "pseudogene" not in f.qualifiers
93 ]
94 none_gene_features = [
95 f
96 for f in record.features
97 if f.type != "gene" and "pseudogene" not in f.qualifiers
98 ]
100 gene_id_to_gene_obj: dict[str, Gene] = {}
101 for gene_feature in gene_features:
102 """
103 gene_feature = type: gene, location: [26244:26472](+), qualifiers: Key: gene, Value: ['E']
104 """
105 gene_symbol = _determine_symbol(gene_feature)
106 gene_accession = _determine_gene_accession(gene_feature)
107 if gene_accession is None:
108 raise ValueError(
109 f"No gene symbol found for feature at {gene_feature.location}."
110 )
112 gene_type = determine_gene_type(none_gene_features, gene_symbol)
113 # TODO : check if join/multiple orders
114 gene_obj = _put_gene_from_feature(
115 gene_feature,
116 record.seq,
117 replicon,
118 gene_type,
119 gene_accession,
120 gene_symbol,
121 )
122 gene_id_to_gene_obj[gene_accession] = gene_obj
123 _create_gene_segments(gene_feature, gene_obj)
125 # features with CDS qualifier
126 gene_id_to_cds_obj: dict[str, list[CDS]] = {}
127 cds_features = [
128 f
129 for f in record.features
130 if f.type == "CDS" and "pseudogene" not in f.qualifiers
131 ]
132 for cds_feature in cds_features:
133 # TODO add note to cds table (e.g. ORF1)
134 gene_accession = _determine_gene_accession(cds_feature)
135 if gene_accession is None:
136 raise ValueError("No gene accession found.")
137 gene = gene_id_to_gene_obj.get(gene_accession, None)
138 if gene is None:
139 raise ValueError("No gene object found for CDS.")
140 cds = _put_cds_from_feature(cds_feature, gene)
141 # different cds in one gene
142 if gene_accession not in gene_id_to_cds_obj:
143 gene_id_to_cds_obj[gene_accession] = []
144 gene_id_to_cds_obj[gene_accession].append(cds)
145 else:
146 gene_id_to_cds_obj[gene_accession].append(cds)
147 _create_cds_segments(cds_feature, cds)
149 # features with peptide qualifier (e.g. HIV)
150 peptide_features = [
151 f for f in record.features if f.type in ["mat_peptide", "sig_peptide"]
152 ]
153 for peptide_feature in peptide_features:
154 gene_accession = _determine_gene_accession(peptide_feature)
155 if gene_accession is None:
156 raise ValueError("No gene symbol found.")
157 cds_objects = gene_id_to_cds_obj.get(gene_accession, None)
158 cds_segments = CDSSegment.objects.filter(
159 cds__in=[cds.pk for cds in cds_objects]
160 )
161 if cds is None:
162 raise ValueError("No gene found for peptide.")
163 elif len(cds_objects) > 1:
164 raise ValueError(
165 f"Multiple CDS objects found for gene symbol {gene_accession}. Can't assign mat_peptides to CDS."
166 )
167 elif len(cds_objects) == 1:
168 cds = cds_objects[0]
169 peptide = _put_peptide_from_feature(peptide_feature, cds)
170 _create_peptide_segments(peptide_feature, peptide, cds_segments)
172 return records
175def determine_gene_type(
176 none_gene_features: list[SeqFeature.SeqFeature], gene_symbol: str
177) -> Gene.GeneTypes | None:
178 types = [
179 Gene.GeneTypes(feature.type)
180 for feature in none_gene_features
181 if feature.type in Gene.GeneTypes.values
182 and feature.qualifiers.get("gene", [None])[0] == gene_symbol
183 ]
184 if len(set(types)) > 1:
185 raise ValueError(f"Multiple gene types found for {gene_symbol}: {types}")
186 return types[0] if types else None
189def _process_segments(
190 feat_location_parts: list[SeqFeature.FeatureLocation | SeqFeature.CompoundLocation],
191 include_strand: bool = False,
192) -> list[dict[str, int]]:
193 """
194 Process the genomic regions (segments) of a feature.
196 Args:
197 feat_location_parts (List[Union[FeatureLocation, CompoundLocation]]): List of feature location parts.
198 cds (bool): A flag indicating whether the segment corresponds to a coding sequence.
199 Default is False.
201 Returns:
202 segments (List[List[int]]): A list of processed segments. Each segment is represented
203 as a list of integers [start, end, strand, index].
204 """
205 segments = []
206 for i, segment in enumerate(feat_location_parts, 1):
207 segment_data = {
208 "start": int(segment.start),
209 "end": int(segment.end),
210 "order": i,
211 }
212 if include_strand:
213 segment_data["forward_strand"] = True if segment.strand == 1 else False
214 segments.append(segment_data)
215 return segments
218def _validate_segment_lengths(parts, accession):
219 parts = _process_segments(parts, include_strand=True)
220 if sum([abs(x["end"] - x["start"]) for x in parts]) % 3 != 0:
221 raise ValueError(f"The length of cds '{accession}' is not a multiple of 3.")
224def _determine_gene_accession(feature: SeqFeature.SeqFeature) -> str | None:
225 if feature.id != "<unknown id>":
226 return feature.id
227 elif "locus_tag" in feature.qualifiers:
228 return feature.qualifiers["locus_tag"][0]
229 elif "gene" in feature.qualifiers:
230 return feature.qualifiers["gene"][0]
231 raise ValueError("No qualifier for gene accession found.")
234def _determine_cds_accession(feature: SeqFeature.SeqFeature) -> str | None:
235 if feature.id != "<unknown id>":
236 return feature.id
237 elif "protein_id" in feature.qualifiers:
238 return feature.qualifiers["protein_id"][0]
239 elif "locus_tag" in feature.qualifiers:
240 return feature.qualifiers["locus_tag"][0]
241 elif "gene" in feature.qualifiers:
242 return feature.qualifiers["gene"][0]
243 raise ValueError("No qualifier for cds accession found.")
246def _determine_symbol(feature: SeqFeature.SeqFeature) -> str | None:
247 if "gene" in feature.qualifiers:
248 return feature.qualifiers["gene"][0]
249 elif "locus_tag" in feature.qualifiers:
250 return feature.qualifiers["locus_tag"][0]
251 raise ValueError("No qualifier for gene symbol found.")
254def _put_gene_from_feature(
255 feature: SeqFeature.SeqFeature,
256 replicon_seq: str,
257 replicon: Replicon,
258 gene_type: Gene.GeneTypes,
259 gene_accession: str,
260 gene_symbol: str,
261) -> Gene:
262 """
263 Processes a gene feature from a gene bank file and updates or creates a corresponding Gene object.
265 Args:
266 feature (SeqFeature.SeqFeature): The gene feature containing location and qualifiers.
267 e.g. type: gene, location: [26244:26472](+), qualifiers: Key: gene, Value: ['E']
268 replicon_seq (str): The full nucleotide sequence of the replicon.
269 replicon (Replicon): The associated Replicon object.
270 gene_type (Gene.GeneTypes): The type of gene (e.g. CDS, rRNA, etc.).
272 Returns:
273 Gene: The updated or newly created Gene object.
274 gene_accession: filled with first existing tag in this order [locus_tag == systematic gene name, gene, protein_id]
275 gene_symbol: filled with first existing tag in this order [gene, locus_tag]
276 gene_sequence: complete sequence of gene
277 gene_description: gene_synonym
279 Raises:
280 ValueError: If the feature has no location information.
281 """
282 if feature.location is None:
283 raise ValueError("No location information found for gene feature.")
284 gene_base_data = {
285 "start": int(feature.location.start),
286 "end": int(feature.location.end),
287 "forward_strand": True if feature.location.strand == 1 else False,
288 "replicon": replicon.pk,
289 "type": gene_type,
290 "accession": gene_accession,
291 "symbol": gene_symbol,
292 }
293 gene_update_data = {}
294 gene_update_data["sequence"] = str(feature.extract(replicon_seq))
295 gene_update_data["description"] = feature.qualifiers.get("gene_synonym", [""])[0]
297 gene = find_or_create(gene_base_data, Gene, GeneSerializer)
298 for attr_name, value in gene_update_data.items():
299 setattr(gene, attr_name, value)
301 return GeneSerializer(gene).update(gene, gene_update_data)
304def _put_cds_from_feature(feature: SeqFeature.SeqFeature, gene: Gene) -> CDS:
305 """
306 Processes a cds feature from a gene bank file and updates or creates a corresponding CDS object.
307 Location information stored in related cds_segment table
309 Args:
310 feature (SeqFeature.SeqFeature): The cds feature containing location and qualifiers.
311 gene (Gene): The associated gene object.
313 Returns:
314 CDS: The updated or newly created CDS object.
315 cds_accession: filled with first existing tag in this order [protein_id, locus_tag, gene]
316 cds_sequence: complete aa-sequence of cds
317 cds_descritption: product tag = protein name
319 Raises:
320 ValueError: If the feature has no location information.
321 """
322 if feature.location is None:
323 raise ValueError("No location information found for CDS feature.")
324 cds_update_data = {}
325 cds_update_data["accession"] = _determine_cds_accession(feature)
326 _validate_segment_lengths(feature.location.parts, cds_update_data["accession"])
327 cds_update_data["sequence"] = feature.qualifiers.get("translation", [""])[0]
328 cds_update_data["description"] = feature.qualifiers.get("product", [""])[0]
329 cds = find_or_create(
330 {"gene": gene.pk, "accession": cds_update_data["accession"]}, CDS, CDSSerializer
331 )
333 for attr_name, value in cds_update_data.items():
334 setattr(cds, attr_name, value)
335 return CDSSerializer(cds).update(cds, cds_update_data)
338def _put_peptide_from_feature(feature: SeqFeature.SeqFeature, cds: CDS) -> Peptide:
339 """
340 Processes a peptide feature from a gene bank file and updates or creates a corresponding Peptide object.
341 Location information stored in related peptide_segment table
343 Args:
344 feature (SeqFeature.SeqFeature): The cds feature containing location and qualifiers.
345 cds (CDS): The associated CDS object.
347 Returns:
348 Peptide: The updated or newly created peptide object.
349 peptide_type: mat_peptide or sig_peptide
350 peptide_descritption: product tag = protein name
352 Raises:
353 ValueError: If the feature has no location information.
354 """
355 if feature.location is None:
356 raise ValueError("No location information found for peptide feature.")
357 peptide_base_data = {
358 "cds": cds.pk,
359 "type": feature.type,
360 "description": feature.qualifiers.get("product", [""])[0],
361 }
362 peptide_update_data = {}
363 peptide = find_or_create(peptide_base_data, Peptide, PeptideSerializer)
364 for attr_name, value in peptide_update_data.items():
365 setattr(peptide, attr_name, value)
366 return PeptideSerializer(peptide).update(peptide, peptide_update_data)
369def parse_date(date_string: str) -> datetime.date:
370 date_formats = ["%Y-%m", "%d-%b-%Y", "%Y", "%Y-%m-%d", "%Y-%b-%d"]
371 for date_format in date_formats:
372 try:
373 return datetime.strptime(date_string, date_format).date()
374 except ValueError:
375 continue
376 # Handle custom format "Dec-2019" without relying on locale (problems can occur in some settings)
377 try:
378 month_str, year = date_string.split("-")
379 month = list(calendar.month_abbr).index(month_str[:3])
380 return datetime(year=int(year), month=month, day=1).date()
381 except (ValueError, IndexError):
382 pass
384 raise ValueError(f"Date string '{date_string}' does not match any known formats.")
387def _put_reference_from_record(
388 record: SeqRecord.SeqRecord, translation_id: int, file_path: str
389) -> Reference:
390 source = None
391 for feature in record.features:
392 if feature.type == "source":
393 source = feature
394 break
395 if source is None:
396 raise Exception("No source feature found.")
397 if "db_xref" in source.qualifiers:
398 if ref := Reference.objects.filter(
399 db_xref=source.qualifiers["db_xref"]
400 ).first():
401 return ref
402 reference = {
403 "accession": f"{record.name}.{record.annotations['sequence_version']}",
404 "description": record.description,
405 "organism": record.annotations["organism"],
406 "translation_id": translation_id,
407 "name": str(file_path),
408 }
409 for attr_name in [
410 "mol_type",
411 "isolate",
412 "host",
413 "db_xref",
414 "country",
415 "collection_date",
416 ]:
417 if attr_name in source.qualifiers:
418 if attr_name == "collection_date":
419 date_string = source.qualifiers[attr_name][0]
420 reference[attr_name] = parse_date(date_string).strftime("%Y-%m-%d")
421 else:
422 reference[attr_name] = source.qualifiers[attr_name][0]
423 else:
424 reference[attr_name] = None
425 try:
426 return Reference.objects.get(db_xref=reference["db_xref"])
427 except Reference.DoesNotExist:
428 serializer = ReferenceSerializer(data=reference)
429 serializer.is_valid(raise_exception=True)
430 return serializer.save()
433def _create_gene_segments(feature: SeqFeature.SeqFeature, gene: Gene):
434 for elempart in _process_segments(feature.location.parts, include_strand=True):
435 elempart_data = {
436 "gene": gene.pk,
437 **elempart,
438 }
439 find_or_create(elempart_data, GeneSegment, GeneSegmentSerializer)
442def _create_cds_segments(feature: SeqFeature.SeqFeature, cds: CDS):
443 for elempart in _process_segments(feature.location.parts, include_strand=True):
444 elempart_data = {
445 "cds": cds.pk,
446 **elempart,
447 }
448 find_or_create(elempart_data, CDSSegment, CDSSegmentSerializer)
451def calculate_cds_start_end(
452 seq_feature: SeqFeature.SeqFeature,
453 cds_segments: list[CDSSegment],
454 cds_accession: str | None = None,
455) -> tuple[int, int]:
456 """
457 Calculate the start and end positions of a peptide segment in relation to the CDS amino acid sequence.
459 Args:
460 seq_feature (SeqFeature.SeqFeature): The peptide feature with location information.
461 cds_segments (list[CDSSegment]): List of CDS segments (start, end, strand, order).
463 Returns:
464 tuple[int, int]: Start and end positions of the peptide in the CDS amino acid sequence.
465 """
466 # Flatten the CDS segments into a continuous nucleotide position list
467 cds_nt_positions = []
468 for segment in cds_segments:
469 if segment.forward_strand:
470 cds_nt_positions.extend(range(segment.start, segment.end + 1))
471 else:
472 cds_nt_positions.extend(range(segment.end, segment.start - 1, -1))
474 # Map peptide segment positions to CDS nucleotide positions
475 peptide_start_nt = int(seq_feature.location.start)
476 peptide_end_nt = int(seq_feature.location.end)
478 # Find the peptide's start and end positions in the CDS nucleotide sequence
479 try:
480 start_cds_nt = cds_nt_positions.index(peptide_start_nt) + 1
481 end_cds_nt = cds_nt_positions.index(peptide_end_nt) + 1
482 except ValueError:
483 raise ValueError("Peptide segment is not part of the CDS segments.")
485 # Convert nucleotide positions to amino acid positions
486 start_cds_aa = (start_cds_nt + 2) // 3 # Convert to 1-based AA position
487 end_cds_aa = (end_cds_nt + 2) // 3 # Convert to 1-based AA position
488 if start_cds_aa < 0 or end_cds_aa < 0:
489 raise ValueError(
490 f"Start or end cds position of peptide segment is out of range for CDS {cds_accession}."
491 )
492 if start_cds_aa == end_cds_aa:
493 raise ValueError(
494 f"Start and end position (start_cds, end_cds) of peptide segment are equal for CDS {cds_accession}."
495 )
496 return start_cds_aa, end_cds_aa
499def _create_peptide_segments(
500 seq_feature: SeqFeature.SeqFeature, peptide, cds_segments: list[CDSSegment]
501):
502 """
503 Create peptide segments based on the provided sequence feature and peptide object.
504 Args:
505 seq_feature (SeqFeature.SeqFeature): The sequence feature containing location information.
506 peptide (Peptide): The associated peptide object.
507 cds_segments (list[CDSSegment]): List of CDS segments (start, end, strand, order) of associated CDS
508 """
509 for elempart in _process_segments(seq_feature.location.parts):
510 start_cds_aa, end_cds_aa = calculate_cds_start_end(
511 seq_feature,
512 cds_segments,
513 peptide.cds.accession,
514 )
515 elempart_data = {
516 "peptide": peptide.pk,
517 "start_cds": start_cds_aa,
518 "end_cds": end_cds_aa,
519 **elempart,
520 }
521 find_or_create(elempart_data, PeptideSegment, PeptideSegmentSerializer)
524def _temp_save_file(uploaded_file: InMemoryUploadedFile):
525 # Create the directory path
526 directory_path = pathlib.Path(SONAR_DATA_ENTRY_FOLDER) / "gbks"
527 directory_path.mkdir(exist_ok=True)
528 file_path = directory_path / uploaded_file.name
529 with open(file_path, "wb") as f:
530 f.write(uploaded_file.read())
531 return file_path