Coverage for django_napse/simulations/models/datasets/dataset.py: 91%

1from datetime import timedelta 

2 

3import pandas as pd 

4from binance.helpers import interval_to_milliseconds 

5from django.db import models 

6 

7from django_napse.simulations.models.datasets.managers.dataset import DataSetManager 

8from django_napse.utils.constants import DOWNLOAD_STATUS 

9from django_napse.utils.errors import DataSetError 

10 

11 

12class DataSet(models.Model): 

13 controller = models.OneToOneField("django_napse_core.Controller", on_delete=models.CASCADE, related_name="dataset") 

14 start_date = models.DateTimeField(null=True, blank=True) 

15 end_date = models.DateTimeField(null=True, blank=True) 

16 

17 last_update = models.DateTimeField(auto_now_add=True) 

18 status = models.CharField(max_length=12, default=DOWNLOAD_STATUS.IDLE) 

19 completion = models.FloatField(default=0.0) 

20 eta = models.DurationField(null=True, blank=True) 

21 

22 objects = DataSetManager() 

23 

24 def __str__(self): 

25 return f"DATASET: {self.pk=}" 

26 

27 def save(self, *args, **kwargs): 

28 if self.completion < 0 or self.completion > 100: 

29 error_msg = f"Completion ({self.completion}) not in [0, 100]" 

30 raise DataSetError.InvalidSettings(error_msg) 

31 

32 if self.status not in str(DOWNLOAD_STATUS): 

33 error_msg = f"Status ({self.status}) not in {DOWNLOAD_STATUS}" 

34 raise DataSetError.InvalidSettings(error_msg) 

35 

36 query = Candle.objects.filter(dataset=self).order_by("open_time") 

37 number_of_candles = query.count() 

38 if number_of_candles > 0: 

39 self.start_date = query[0].open_time 

40 self.end_date = query[number_of_candles - 1].open_time + timedelta(milliseconds=interval_to_milliseconds(self.controller.interval) - 1) 

41 super().save(*args, **kwargs) 

42 

43 def info(self, verbose=True, beacon=""): 

44 string = "" 

45 string += f"{beacon}Dataset {self.pk}:\n" 

46 string += f"{beacon}\t{self.controller=}\n" 

47 string += f"{beacon}\t{self.start_date=}\n" 

48 string += f"{beacon}\t{self.end_date=}\n" 

49 string += f"{beacon}\t{self.last_update=}\n" 

50 string += f"{beacon}\t{self.status=}\n" 

51 string += f"{beacon}\t{self.completion=}\n" 

52 string += f"{beacon}\t{self.eta=}\n" 

53 string += f"{beacon}Candles:\n" 

54 candles = self.candles.all().order_by("open_time") 

55 if candles.count() > 10: 

56 for candle in candles[:5]: 

57 string += f"{beacon}\tO:{candle.open_time}\t H: {candle.high}\t L: {candle.close}\t C: {candle.low}\t V: {candle.volume}\n" 

58 string += f"{beacon}\t...\n" 

59 for candle in candles[candles.count() - 5 :]: 

60 string += f"{beacon}\tO:{candle.open_time}\t H: {candle.high}\t L: {candle.close}\t C: {candle.low}\t V: {candle.volume}\n" 

61 string += f"{beacon}\t({candles.count()} candles)\n" 

62 else: 

63 for candle in candles: 

64 string += f"{beacon}\tO:{candle.open_time}\t H: {candle.high}\t L: {candle.close}\t C: {candle.low}\t V: {candle.volume}\n" 

65 

66 if verbose: 

67 print(string) 

68 return string 

69 

70 def create_candles(self, candles: list): 

71 """Save a list of candle Objects into the database.""" 

72 Candle.objects.bulk_create(candles) 

73 

74 def set_downloading(self): 

75 """Set the dataset status to downloading. 

76 

77 Raises: 

78 ValueError: If the dataset isn't in IDLE status. 

79 """ 

80 if self.status == DOWNLOAD_STATUS.IDLE: 

81 self.status = DOWNLOAD_STATUS.DOWNLOADING 

82 self.completion = 0 

83 self.save() 

84 else: 

85 error_msg = "Dataset is already downloading." 

86 raise DataSetError.InvalidSettings(error_msg) 

87 

88 def set_idle(self): 

89 """Set the dataset status to idle.""" 

90 if self.status == DOWNLOAD_STATUS.DOWNLOADING: 

91 self.status = DOWNLOAD_STATUS.IDLE 

92 self.save() 

93 else: 

94 error_msg = "Dataset is not downloading." 

95 raise DataSetError.InvalidSettings(error_msg) 

96 

97 def is_finished(self): 

98 return self.status == DOWNLOAD_STATUS.IDLE and self.completion == 100 

99 

100 def to_dataframe(self, start_date=None, end_date=None): 

101 start_date = start_date or self.start_date 

102 end_date = end_date or self.end_date 

103 candles = self.candles.filter(open_time__gte=start_date, open_time__lt=end_date).order_by("open_time") 

104 return pd.DataFrame( 

105 data={ 

106 "open_time": [candle.open_time for candle in candles], 

107 "open": [candle.open for candle in candles], 

108 "high": [candle.high for candle in candles], 

109 "low": [candle.low for candle in candles], 

110 "close": [candle.close for candle in candles], 

111 "volume": [candle.volume for candle in candles], 

112 }, 

113 ) 

114 

115 

116class Candle(models.Model): 

117 dataset = models.ForeignKey("DataSet", on_delete=models.CASCADE, related_name="candles") 

118 open_time = models.DateTimeField() 

119 open = models.FloatField() 

120 high = models.FloatField() 

121 low = models.FloatField() 

122 close = models.FloatField() 

123 volume = models.FloatField() 

124 

125 class Meta: 

126 unique_together = ("dataset", "open_time") 

127 

128 def __str__(self): 

129 return f"CANDLE: {self.pk=}" 

130 

131 def info(self, verbose=True, beacon=""): 

132 string = "" 

133 string += f"{beacon}Candle {self.pk}:\n" 

134 string += f"{beacon}\t{self.dataset=}\n" 

135 string += f"{beacon}\t{self.open_time=}\n" 

136 string += f"{beacon}\t{self.open=}\n" 

137 string += f"{beacon}\t{self.high=}\n" 

138 string += f"{beacon}\t{self.low=}\n" 

139 string += f"{beacon}\t{self.close=}\n" 

140 string += f"{beacon}\t{self.volume=}\n" 

141 

142 if verbose: 

143 print(string) 

144 return string 

145 

146 def to_dict(self): 

147 return { 

148 "controller": self.dataset.controller, 

149 "open_time": self.open_time, 

150 "open": self.open, 

151 "high": self.high, 

152 "low": self.low, 

153 "close": self.close, 

154 "volume": self.volume, 

155 "extra": {}, 

156 } 

157 

158 

159class DataSetQueue(models.Model): 

160 controller = models.ForeignKey("django_napse_core.Controller", on_delete=models.CASCADE, related_name="dataset_queues") 

161 start_date = models.DateTimeField() 

162 end_date = models.DateTimeField() 

163 

164 created_at = models.DateTimeField(auto_now_add=True) 

165 

166 def __str__(self) -> str: 

167 return f"DATASET_QUEUE: {self.pk=}" 

168 

169 def get_dataset(self) -> DataSet: 

170 try: 

171 return DataSet.objects.get(controller=self.controller) 

172 except DataSet.DoesNotExist: 

173 return None 

174 return None