Coverage for middle_layer/common/application_layer/orm_repositories/orm_types.py: 78.16%

1# Copyright 2024 Associated Universities, Inc. 

2# 

3# This file is part of Telescope Time Allocation Tools (TTAT). 

4# 

5# TTAT is free software: you can redistribute it and/or modify 

6# it under the terms of the GNU General Public License as published by 

7# the Free Software Foundation, either version 3 of the License, or 

8# any later version. 

9# 

10# TTAT is distributed in the hope that it will be useful, 

11# but WITHOUT ANY WARRANTY; without even the implied warranty of 

12# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

13# GNU General Public License for more details. 

14# 

15# You should have received a copy of the GNU General Public License 

16# along with TTAT. If not, see <https://www.gnu.org/licenses/>. 

17# 

18from typing import Any, Type, override 

19 

20import sqlalchemy.types as types 

21from astropy.units import Quantity, Unit, degree, s 

22 

23from common.domain_layer import JSON 

24 

25 

26class QuantitySeconds(types.TypeDecorator): 

27 """converts an astropy Quantity in seconds into Numeric for storage.""" 

28 

29 impl = types.Numeric 

30 

31 cache_ok = True 

32 

33 @override 

34 def process_bind_param(self, value: Quantity["time"], dialect): 

35 return value.to_value(s) 

36 

37 @override 

38 def process_result_value(self, value, dialect) -> Quantity[s]: 

39 # convert the numeric from the db into Quantity in seconds: 

40 return Quantity(value, s) 

41 

42 @property 

43 @override 

44 def python_type(self) -> type: 

45 return Quantity 

46 

47 

48class QuantityDegrees(types.TypeDecorator): 

49 """converts an astropy Quantity in degrees into Numeric for storage.""" 

50 

51 impl = types.Numeric 

52 

53 cache_ok = True 

54 

55 @override 

56 def process_bind_param(self, value, dialect): 

57 # check that the input is a Quantity in seconds: 

58 if not (isinstance(value, Quantity) and value.unit == Unit("degree")): 

59 try: 

60 value = value * Unit("degree") 

61 except ValueError: 

62 raise TypeError(f"value {value} must be a Quantity in degrees.") 

63 # return just the numeric part of the value 

64 return value.value 

65 

66 @override 

67 def process_result_value(self, value, dialect): 

68 # convert the numeric from the db into Quantity in seconds: 

69 return Quantity(value, degree) 

70 

71 @property 

72 @override 

73 def python_type(self) -> Type[Any]: 

74 return Quantity 

75 

76 

77class JSONList(types.TypeDecorator): 

78 """Convert a list of JSON values into a single JSON blob for storage, and vice versa""" 

79 

80 impl = types.JSON(none_as_null=True) 

81 cache_ok = True 

82 

83 @override 

84 def process_bind_param(self, value: list[JSON], dialect) -> JSON: 

85 return value 

86 

87 @override 

88 def process_result_value(self, value: JSON, dialect) -> list[JSON]: 

89 return value 

90 

91 

92class CSV(types.TypeDecorator): 

93 """Convert a list of strings into a single String for storage, and vice versa""" 

94 

95 SEPARATOR = "," 

96 impl = types.String 

97 cache_ok = True 

98 

99 @override 

100 def process_bind_param(self, value: list[str], dialect) -> str: 

101 return self.SEPARATOR.join(value) 

102 

103 @override 

104 def process_result_value(self, value: str | None, dialect) -> list[str]: 

105 return value.split(self.SEPARATOR) if value is not None else [] 

106 

107 @property 

108 @override 

109 def python_type(self) -> Type[Any]: 

110 return list[str] 

111 

112 

113class CSVFloat(types.TypeDecorator): 

114 """Convert a list of floats into a single String for storage, and vice versa""" 

115 

116 SEPARATOR = "," 

117 impl = types.String 

118 cache_ok = True 

119 

120 @override 

121 def process_bind_param(self, value: list[float], dialect) -> str: 

122 return self.SEPARATOR.join([str(i) for i in value]) 

123 

124 @override 

125 def process_result_value(self, value: str | None, dialect) -> list[float]: 

126 return [float(x) for x in value.split(self.SEPARATOR)] if value is not None else [] 

127 

128 @property 

129 @override 

130 def python_type(self) -> Type[Any]: 

131 return list[float] 

132 

133 

134class CSVInt(types.TypeDecorator): 

135 """Convert a list of floats into a single String for storage, and vice versa""" 

136 

137 "It would be great to combine CSV, CSVFloat, and CSVInt into one type, if possible." 

138 

139 SEPARATOR = "," 

140 impl = types.String 

141 cache_ok = True 

142 

143 @override 

144 def process_bind_param(self, value: list[int], dialect) -> str: 

145 return self.SEPARATOR.join([str(i) for i in value]) 

146 

147 @override 

148 def process_result_value(self, value: str | None, dialect) -> list[int]: 

149 return [int(x) for x in value.split(self.SEPARATOR)] if value is not None else [] 

150 

151 @property 

152 @override 

153 def python_type(self) -> Type[Any]: 

154 return list[int]