Coverage for rest_api/management/commands/import_lineage.py: 29%

1import os 

2import shutil 

3from tempfile import mkdtemp 

4from typing import List 

5from typing import Optional 

6 

7from django.core.management.base import BaseCommand 

8from django.core.management.base import CommandError 

9from django.db import transaction 

10import pandas as pd 

11 

12from rest_api.models import Lineage 

13from rest_api.models import Reference 

14 

15 

16class LineageImport: 

17 """ 

18 sonarLinmgr class is used to manage the lineages. 

19 

20 Attributes: 

21 _tmpdir (str): Temporary directory. 

22 _linurl (str): URL of the lineages file. 

23 _aliurl (str): URL of the alias file. 

24 lineage_file (str): Local path of the lineages file. 

25 alias_file (str): Local path of the alias file. 

26 """ 

27 

28 def __init__(self, tmpdir: Optional[str] = None): 

29 """ 

30 sonarLinmgr Constructor. 

31 

32 Args: 

33 tmpdir (str): Temporary directory. 

34 """ 

35 self._tmpdir = mkdtemp(prefix=".tmp_sonarLinmgr_") 

36 self.lineage_file = os.path.join(self._tmpdir, "lineages.csv") 

37 

38 def __enter__(self): 

39 return self 

40 

41 def __exit__(self, exc_type, exc_value, exc_traceback): 

42 shutil.rmtree(self._tmpdir) 

43 

44 def set_file(self, file: str | None = None) -> None: 

45 if file: 

46 self.lineage_file = file 

47 return 

48 # Can we remove this? playwright-e2e is failing without 

49 self.lineage_file = "lineage-test-data/lineages_test.tsv" 

50 

51 def process_lineage_data(self, reference: Reference): 

52 """ 

53 Process the lineage data. 

54 

55 Args: 

56 reference (Reference): The reference all imported lineages belong to. 

57 """ 

58 tsv_data = pd.read_csv(self.lineage_file, sep="\t") 

59 

60 # Use a single dictionary to track all lineage objects 

61 all_lineages: dict[str, Lineage] = {} 

62 

63 # First pass: Create all lineage objects 

64 for lineage, sublineages in tsv_data.itertuples(index=False): 

65 # Clean and validate lineage name 

66 lineage = str(lineage).strip() 

67 if len(lineage) > 50: 

68 raise ValueError( 

69 f"Lineage name too long ({len(lineage)} chars): {lineage}" 

70 ) 

71 

72 # Ensure the lineage is added even if it has no children 

73 if lineage not in all_lineages: 

74 all_lineages[lineage] = Lineage(name=lineage, reference=reference) 

75 

76 # Process sublineages if they exist and create them if needed 

77 if sublineages != "none": 

78 values = sublineages.split(",") 

79 for val in values: 

80 val = val.strip() # Remove whitespace 

81 if len(val) > 50: 

82 raise ValueError( 

83 f"Sublineage name too long " f"({len(val)} chars): {val}" 

84 ) 

85 if val not in all_lineages: 

86 all_lineages[val] = Lineage(name=val, reference=reference) 

87 

88 # Second pass: Set parent relationships 

89 for lineage, sublineages in tsv_data.itertuples(index=False): 

90 lineage = str(lineage).strip() 

91 if sublineages != "none": 

92 values = sublineages.split(",") 

93 for val in values: 

94 val = val.strip() 

95 # Set parent only if not set (first occurrence wins) 

96 if all_lineages[val].parent is None: 

97 all_lineages[val].parent = all_lineages[lineage] 

98 

99 # Save all lineages to the database 

100 # Separate into parents (no parent set) and children (parent set) 

101 parents_to_save = [ 

102 lineage_obj 

103 for lineage_obj in all_lineages.values() 

104 if lineage_obj.parent is None 

105 ] 

106 children_to_save = [ 

107 lineage_obj 

108 for lineage_obj in all_lineages.values() 

109 if lineage_obj.parent is not None 

110 ] 

111 

112 with transaction.atomic(): 

113 # Save parents first so they have IDs 

114 for parent in parents_to_save: 

115 parent.save() 

116 # Then save children 

117 for child in children_to_save: 

118 child.save() 

119 

120 def update_lineage_data(self, lineages: str, reference: Reference) -> List[Lineage]: 

121 """ 

122 Update the lineage data. 

123 

124 Returns: 

125 pd.DataFrame: The dataframe with updated data. 

126 """ 

127 if lineages: 

128 self.lineage_file = lineages 

129 else: 

130 self.set_file() 

131 

132 df = self.process_lineage_data(reference) 

133 return df 

134 

135 

136class Command(BaseCommand): 

137 help = "import lineage tsv" 

138 

139 def add_arguments(self, parser): 

140 parser.add_argument( 

141 "--lineages", 

142 help="lineages.tsv file", 

143 type=str, 

144 default=None, 

145 ) 

146 parser.add_argument( 

147 "--reference", 

148 help="accession of the reference the lineages belong to", 

149 type=str, 

150 required=True, 

151 ) 

152 

153 def handle(self, *args, **kwargs): 

154 try: 

155 reference = Reference.objects.get(accession=kwargs["reference"]) 

156 except Reference.DoesNotExist: 

157 raise CommandError( 

158 f"Reference with accession '{kwargs['reference']}' not found. " 

159 "Import the reference before importing its lineages." 

160 ) 

161 # Only replace the lineages of this reference, keep the others intact. 

162 Lineage.objects.filter(reference=reference).delete() 

163 with LineageImport() as lineage_manager: 

164 lineages = lineage_manager.update_lineage_data( 

165 kwargs["lineages"], reference 

166 ) 

167 print("--Done--")