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

1import calendar 

2from datetime import datetime 

3import os 

4import pathlib 

5 

6from Bio import SeqFeature 

7from Bio import SeqIO 

8from Bio import SeqRecord 

9from django.core.files.uploadedfile import InMemoryUploadedFile 

10from django.db import transaction 

11 

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 

30 

31 

32def import_gbk_file(uploaded_files: list[InMemoryUploadedFile], translation_id: int): 

33 """ 

34 Import GenBank file. or Import multiple GenBank files as segments. 

35 

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"))) 

46 

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"] 

78 

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] 

86 

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 ] 

99 

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 ) 

111 

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) 

124 

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) 

148 

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) 

171 

172 return records 

173 

174 

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 

187 

188 

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. 

195 

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. 

200 

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 

216 

217 

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.") 

222 

223 

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.") 

232 

233 

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.") 

244 

245 

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.") 

252 

253 

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. 

264 

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.). 

271 

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 

278 

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] 

296 

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) 

300 

301 return GeneSerializer(gene).update(gene, gene_update_data) 

302 

303 

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 

308 

309 Args: 

310 feature (SeqFeature.SeqFeature): The cds feature containing location and qualifiers. 

311 gene (Gene): The associated gene object. 

312 

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 

318 

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 ) 

332 

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) 

336 

337 

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 

342 

343 Args: 

344 feature (SeqFeature.SeqFeature): The cds feature containing location and qualifiers. 

345 cds (CDS): The associated CDS object. 

346 

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 

351 

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) 

367 

368 

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 

383 

384 raise ValueError(f"Date string '{date_string}' does not match any known formats.") 

385 

386 

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() 

431 

432 

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) 

440 

441 

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) 

449 

450 

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. 

458 

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). 

462 

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)) 

473 

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) 

477 

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.") 

484 

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 

497 

498 

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) 

522 

523 

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