oafuncs 0.0.97.16__py3-none-any.whl → 0.0.97.17__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
oafuncs/oa_data.py CHANGED
@@ -16,278 +16,217 @@ Python Version: 3.11
16
16
  import itertools
17
17
  import multiprocessing as mp
18
18
  from concurrent.futures import ThreadPoolExecutor
19
-
19
+ from typing import Any, List, Union
20
20
 
21
21
  import numpy as np
22
22
  import salem
23
23
  import xarray as xr
24
+ from rich import print
24
25
  from scipy.interpolate import griddata, interp1d
25
26
 
26
27
  __all__ = ["interp_along_dim", "interp_2d", "ensure_list", "mask_shapefile"]
27
28
 
28
29
 
29
- def ensure_list(input_data):
30
+ def ensure_list(input_value: Any) -> List[str]:
30
31
  """
31
- Ensures that the input is converted into a list.
32
+ Ensure the input is converted into a list.
32
33
 
33
- If the input is already a list, it returns it directly.
34
- If the input is a string, it wraps it in a list and returns.
35
- For other types of input, it converts them to a string and then wraps in a list.
34
+ Args:
35
+ input_value (Any): The input which can be a list, a string, or any other type.
36
36
 
37
- :param input_data: The input which can be a list, a string, or any other type.
38
- :return: A list containing the input or the string representation of the input.
37
+ Returns:
38
+ List[str]: A list containing the input or the string representation of the input.
39
39
  """
40
- if isinstance(input_data, list):
41
- return input_data
42
- elif isinstance(input_data, str):
43
- return [input_data]
40
+ if isinstance(input_value, list):
41
+ return input_value
42
+ elif isinstance(input_value, str):
43
+ return [input_value]
44
44
  else:
45
- # For non-list and non-string inputs, convert to string and wrap in a list
46
- return [str(input_data)]
45
+ return [str(input_value)]
47
46
 
48
47
 
49
- def interp_along_dim(tgt_coords, src_coords, src_data, axis=-1, interp_method="linear", extrap_method="linear"):
48
+ def interp_along_dim(
49
+ target_coordinates: np.ndarray,
50
+ source_coordinates: Union[np.ndarray, List[float]],
51
+ source_data: np.ndarray,
52
+ interpolation_axis: int = -1,
53
+ interpolation_method: str = "linear",
54
+ extrapolation_method: str = "linear",
55
+ ) -> np.ndarray:
50
56
  """
51
- 在指定维度上执行插值和外插操作。
52
-
53
- Parameters:
54
- -----------
55
- tgt_coords: 1d array
56
- 目标坐标点数组,必须是一维数组。
57
-
58
- src_coords: 1d or nd array
59
- 源坐标点数组。可以是一维数组(将被广播到与src_data匹配的形状)或与src_data相同形状的多维数组。
60
-
61
- src_data: nd array
62
- 源数据数组,包含要从src_coords插值到tgt_coords的现象值。
63
-
64
- axis: int (default -1)
65
- 要在src_data上执行插值的轴。默认为最后一个轴。
66
-
67
- interp_method: str (default "linear")
68
- 核心插值方法。
69
- 可选值包括:
70
- - "linear": 线性插值(默认)
71
- - "nearest": 最近邻插值
72
- - "zero": 零阶插值
73
- - "slinear": 样条一阶插值
74
- - "quadratic": 二阶插值
75
- - "cubic": 三阶插值
76
- - "previous": 使用前一个点的值
77
- - "next": 使用后一个点的值
78
- 更多选项参考scipy.interpolate.interp1d的kind参数。
79
-
80
- extrap_method: str (default "linear")
81
- 核心外插方法,用于处理超出源坐标范围的目标坐标点。
82
- 支持与interp_method相同的选项:
83
- - "linear": 线性外插(默认)
84
- - "nearest": 最近邻外插
85
- - "zero": 零阶外插
86
- - "slinear": 样条一阶外插
87
- - "quadratic": 二阶外插
88
- - "cubic": 三阶外插
89
- - "previous": 使用最近的前一个点的值
90
- - "next": 使用最近的后一个点的值
57
+ Perform interpolation and extrapolation along a specified dimension.
58
+
59
+ Args:
60
+ target_coordinates (np.ndarray): 1D array of target coordinate points.
61
+ source_coordinates (Union[np.ndarray, List[float]]): Source coordinate points (1D or ND array).
62
+ source_data (np.ndarray): Source data array to interpolate.
63
+ interpolation_axis (int, optional): Axis to perform interpolation on. Defaults to -1.
64
+ interpolation_method (str, optional): Interpolation method. Defaults to "linear".
65
+ extrapolation_method (str, optional): Extrapolation method. Defaults to "linear".
91
66
 
92
67
  Returns:
93
- --------
94
- array
95
- 插值后的数据数组,形状将与src_data相同,但在axis轴上长度为len(tgt_coords)。
68
+ np.ndarray: Interpolated data array.
69
+
70
+ Raises:
71
+ ValueError: If input dimensions or shapes are invalid.
96
72
 
97
73
  Examples:
98
- ---------
99
- 1D插值示例:
100
- >>> tgt_coords = np.array([1, 2, 3, 4])
101
- >>> src_coords = np.array([0, 1, 2, 3, 4, 5])
102
- >>> src_data = np.array([0, 1, 4, 9, 16, 25])
103
- >>> interp_along_dim(tgt_coords, src_coords, src_data)
104
- array([ 1., 4., 9., 16.])
105
-
106
- 多维插值示例:
107
- >>> src_data = np.array([[0, 1, 4], [10, 20, 30]])
108
- >>> interp_along_dim(np.array([0.5, 1.5]), np.array([0, 1, 2]), src_data, axis=1)
109
- array([[ 0.5, 2.5],
110
- [15. , 25. ]])
74
+ >>> target_coordinates = np.array([1, 2, 3])
75
+ >>> source_coordinates = np.array([0, 1, 2, 3])
76
+ >>> source_data = np.array([10, 20, 30, 40])
77
+ >>> result = interp_along_dim(target_coordinates, source_coordinates, source_data)
78
+ >>> print(result) # Expected output: [20.0, 30.0]
111
79
  """
112
- tgt_coords = np.asarray(tgt_coords)
113
- if tgt_coords.ndim != 1:
114
- raise ValueError("tgt_coords must be a 1d array.")
115
-
116
- src_coords = np.asarray(src_coords)
117
- src_data = np.asarray(src_data)
118
-
119
- # 处理1维的简单情况
120
- if src_data.ndim == 1 and src_coords.ndim == 1:
121
- if len(src_coords) != len(src_data):
122
- raise ValueError("For 1D data, src_coords and src_data must have the same length")
123
-
124
- interpolator = interp1d(src_coords, src_data, kind=interp_method, fill_value="extrapolate", bounds_error=False)
125
- return interpolator(tgt_coords)
126
-
127
- # 多维情况的处理
128
- if src_coords.ndim == 1:
129
- # Expand src_coords to match src_data dimensions along the specified axis
130
- shape = [1] * src_data.ndim
131
- shape[axis] = src_coords.shape[0]
132
- src_coords = np.reshape(src_coords, shape)
133
- src_coords = np.broadcast_to(src_coords, src_data.shape)
134
- elif src_coords.shape != src_data.shape:
135
- raise ValueError("src_coords and src_data must have the same shape.")
136
-
137
- def apply_interp_extrap(arr):
138
- xp = np.moveaxis(src_coords, axis, 0)
139
- # 根据维度正确获取坐标
140
- if xp.ndim > 1:
141
- xp = xp[:, 0] # 多维情况
142
- else:
143
- xp = xp # 1维情况
144
-
145
- arr = np.moveaxis(arr, axis, 0)
146
- interpolator = interp1d(xp, arr, kind=interp_method, fill_value="extrapolate", bounds_error=False)
147
- interpolated = interpolator(tgt_coords)
148
- if extrap_method != interp_method:
149
- mask_extrap = (tgt_coords < xp.min()) | (tgt_coords > xp.max())
80
+ target_coordinates = np.asarray(target_coordinates)
81
+ if target_coordinates.ndim != 1:
82
+ raise ValueError("[red]target_coordinates must be a 1D array.[/red]")
83
+
84
+ source_coordinates = np.asarray(source_coordinates)
85
+ source_data = np.asarray(source_data)
86
+
87
+ if source_data.ndim == 1 and source_coordinates.ndim == 1:
88
+ if len(source_coordinates) != len(source_data):
89
+ raise ValueError("[red]For 1D data, source_coordinates and source_data must have the same length.[/red]")
90
+
91
+ interpolator = interp1d(source_coordinates, source_data, kind=interpolation_method, fill_value="extrapolate", bounds_error=False)
92
+ return interpolator(target_coordinates)
93
+
94
+ if source_coordinates.ndim == 1:
95
+ shape = [1] * source_data.ndim
96
+ shape[interpolation_axis] = source_coordinates.shape[0]
97
+ source_coordinates = np.reshape(source_coordinates, shape)
98
+ source_coordinates = np.broadcast_to(source_coordinates, source_data.shape)
99
+ elif source_coordinates.shape != source_data.shape:
100
+ raise ValueError("[red]source_coordinates and source_data must have the same shape.[/red]")
101
+
102
+ def apply_interp_extrap(arr: np.ndarray) -> np.ndarray:
103
+ xp = np.moveaxis(source_coordinates, interpolation_axis, 0)
104
+ xp = xp[:, 0] if xp.ndim > 1 else xp
105
+ arr = np.moveaxis(arr, interpolation_axis, 0)
106
+ interpolator = interp1d(xp, arr, kind=interpolation_method, fill_value="extrapolate", bounds_error=False)
107
+ interpolated = interpolator(target_coordinates)
108
+ if extrapolation_method != interpolation_method:
109
+ mask_extrap = (target_coordinates < xp.min()) | (target_coordinates > xp.max())
150
110
  if np.any(mask_extrap):
151
- extrap_interpolator = interp1d(xp, arr, kind=extrap_method, fill_value="extrapolate", bounds_error=False)
152
- interpolated[mask_extrap] = extrap_interpolator(tgt_coords[mask_extrap])
153
- return np.moveaxis(interpolated, 0, axis)
154
-
155
- result = np.apply_along_axis(apply_interp_extrap, axis, src_data)
156
-
157
- return result
158
-
159
-
160
- def interp_2d(target_x, target_y, origin_x, origin_y, data, method="linear", parallel=True):
111
+ extrap_interpolator = interp1d(xp, arr, kind=extrapolation_method, fill_value="extrapolate", bounds_error=False)
112
+ interpolated[mask_extrap] = extrap_interpolator(target_coordinates[mask_extrap])
113
+ return np.moveaxis(interpolated, 0, interpolation_axis)
114
+
115
+ return np.apply_along_axis(apply_interp_extrap, interpolation_axis, source_data)
116
+
117
+
118
+ def interp_2d(
119
+ target_x_coordinates: Union[np.ndarray, List[float]],
120
+ target_y_coordinates: Union[np.ndarray, List[float]],
121
+ source_x_coordinates: Union[np.ndarray, List[float]],
122
+ source_y_coordinates: Union[np.ndarray, List[float]],
123
+ source_data: np.ndarray,
124
+ interpolation_method: str = "linear",
125
+ use_parallel: bool = True,
126
+ ) -> np.ndarray:
161
127
  """
162
128
  Perform 2D interpolation on the last two dimensions of a multi-dimensional array.
163
129
 
164
- Parameters:
165
- - target_x (array-like): 1D or 2D array of target grid's x-coordinates.
166
- - target_y (array-like): 1D or 2D array of target grid's y-coordinates.
167
- - origin_x (array-like): 1D or 2D array of original grid's x-coordinates.
168
- - origin_y (array-like): 1D or 2D array of original grid's y-coordinates.
169
- - data (numpy.ndarray): Multi-dimensional array where the last two dimensions correspond to the original grid.
170
- - method (str, optional): Interpolation method, default is 'linear'. Other options include 'nearest', 'cubic', etc.
171
- - parallel (bool, optional): Flag to enable parallel processing. Default is True.
130
+ Args:
131
+ target_x_coordinates (Union[np.ndarray, List[float]]): Target grid's x-coordinates.
132
+ target_y_coordinates (Union[np.ndarray, List[float]]): Target grid's y-coordinates.
133
+ source_x_coordinates (Union[np.ndarray, List[float]]): Original grid's x-coordinates.
134
+ source_y_coordinates (Union[np.ndarray, List[float]]): Original grid's y-coordinates.
135
+ source_data (np.ndarray): Multi-dimensional array with the last two dimensions as spatial.
136
+ interpolation_method (str, optional): Interpolation method. Defaults to "linear".
137
+ use_parallel (bool, optional): Enable parallel processing. Defaults to True.
172
138
 
173
139
  Returns:
174
- - interpolated_data (numpy.ndarray): Interpolated data with the same leading dimensions as the input data, but with the last two dimensions corresponding to the target grid.
140
+ np.ndarray: Interpolated data array.
175
141
 
176
142
  Raises:
177
- - ValueError: If the shape of the data does not match the shape of the origin_x or origin_y grids.
178
-
179
- Usage:
180
- - Interpolate a 2D array:
181
- result = interp_2d(target_x, target_y, origin_x, origin_y, data_2d)
182
- - Interpolate a 3D array (where the last two dimensions are spatial):
183
- result = interp_2d(target_x, target_y, origin_x, origin_y, data_3d)
184
- - Interpolate a 4D array (where the last two dimensions are spatial):
185
- result = interp_2d(target_x, target_y, origin_x, origin_y, data_4d)
143
+ ValueError: If input shapes are invalid.
144
+
145
+ Examples:
146
+ >>> target_x_coordinates = np.array([1, 2, 3])
147
+ >>> target_y_coordinates = np.array([4, 5, 6])
148
+ >>> source_x_coordinates = np.array([7, 8, 9])
149
+ >>> source_y_coordinates = np.array([10, 11, 12])
150
+ >>> source_data = np.random.rand(3, 3)
151
+ >>> result = interp_2d(target_x_coordinates, target_y_coordinates, source_x_coordinates, source_y_coordinates, source_data)
152
+ >>> print(result.shape) # Expected output: (3, 3)
186
153
  """
187
154
 
188
- def interp_single(data_slice, target_points, origin_points, method):
189
- return griddata(origin_points, data_slice.ravel(), target_points, method=method).reshape(target_y.shape)
155
+ def interp_single(data_slice: np.ndarray, target_points: np.ndarray, origin_points: np.ndarray, method: str) -> np.ndarray:
156
+ return griddata(origin_points, data_slice.ravel(), target_points, method=method).reshape(target_y_coordinates.shape)
190
157
 
191
- # 确保目标网格和初始网格都是二维的
192
- if len(target_y.shape) == 1:
193
- target_x, target_y = np.meshgrid(target_x, target_y)
194
- if len(origin_y.shape) == 1:
195
- origin_x, origin_y = np.meshgrid(origin_x, origin_y)
158
+ if len(target_y_coordinates.shape) == 1:
159
+ target_x_coordinates, target_y_coordinates = np.meshgrid(target_x_coordinates, target_y_coordinates)
160
+ if len(source_y_coordinates.shape) == 1:
161
+ source_x_coordinates, source_y_coordinates = np.meshgrid(source_x_coordinates, source_y_coordinates)
196
162
 
197
- # 根据经纬度网格判断输入数据的形状是否匹配
198
- if origin_x.shape != data.shape[-2:] or origin_y.shape != data.shape[-2:]:
199
- raise ValueError("Shape of data does not match shape of origin_x or origin_y.")
163
+ if source_x_coordinates.shape != source_data.shape[-2:] or source_y_coordinates.shape != source_data.shape[-2:]:
164
+ raise ValueError("[red]Shape of source_data does not match shape of source_x_coordinates or source_y_coordinates.[/red]")
200
165
 
201
- # 创建网格和展平数据
202
- target_x, target_y = np.array(target_x), np.array(target_y)
203
- origin_x, origin_y = np.array(origin_x), np.array(origin_y)
204
- target_points = np.column_stack((target_y.ravel(), target_x.ravel()))
205
- origin_points = np.column_stack((origin_y.ravel(), origin_x.ravel()))
166
+ target_points = np.column_stack((np.array(target_y_coordinates).ravel(), np.array(target_x_coordinates).ravel()))
167
+ origin_points = np.column_stack((np.array(source_y_coordinates).ravel(), np.array(source_x_coordinates).ravel()))
206
168
 
207
- # 根据是否并行选择不同的执行方式
208
- if parallel:
169
+ if use_parallel:
209
170
  with ThreadPoolExecutor(max_workers=mp.cpu_count() - 2) as executor:
210
- if len(data.shape) == 2:
211
- interpolated_data = list(executor.map(interp_single, [data], [target_points], [origin_points], [method]))
212
- elif len(data.shape) == 3:
213
- interpolated_data = list(executor.map(interp_single, [data[i] for i in range(data.shape[0])], [target_points] * data.shape[0], [origin_points] * data.shape[0], [method] * data.shape[0]))
214
- elif len(data.shape) == 4:
215
- index_combinations = list(itertools.product(range(data.shape[0]), range(data.shape[1])))
216
- interpolated_data = list(executor.map(interp_single, [data[i, j] for i, j in index_combinations], [target_points] * len(index_combinations), [origin_points] * len(index_combinations), [method] * len(index_combinations)))
217
- interpolated_data = np.array(interpolated_data).reshape(data.shape[0], data.shape[1], *target_y.shape)
171
+ if len(source_data.shape) == 2:
172
+ interpolated_data = list(executor.map(interp_single, [source_data], [target_points], [origin_points], [interpolation_method]))
173
+ elif len(source_data.shape) == 3:
174
+ interpolated_data = list(executor.map(interp_single, [source_data[i] for i in range(source_data.shape[0])], [target_points] * source_data.shape[0], [origin_points] * source_data.shape[0], [interpolation_method] * source_data.shape[0]))
175
+ elif len(source_data.shape) == 4:
176
+ index_combinations = list(itertools.product(range(source_data.shape[0]), range(source_data.shape[1])))
177
+ interpolated_data = list(executor.map(interp_single, [source_data[i, j] for i, j in index_combinations], [target_points] * len(index_combinations), [origin_points] * len(index_combinations), [interpolation_method] * len(index_combinations)))
178
+ interpolated_data = np.array(interpolated_data).reshape(source_data.shape[0], source_data.shape[1], *target_y_coordinates.shape)
218
179
  else:
219
- if len(data.shape) == 2:
220
- interpolated_data = interp_single(data, target_points, origin_points, method)
221
- elif len(data.shape) == 3:
222
- interpolated_data = np.stack([interp_single(data[i], target_points, origin_points, method) for i in range(data.shape[0])])
223
- elif len(data.shape) == 4:
224
- interpolated_data = np.stack([np.stack([interp_single(data[i, j], target_points, origin_points, method) for j in range(data.shape[1])]) for i in range(data.shape[0])])
180
+ if len(source_data.shape) == 2:
181
+ interpolated_data = interp_single(source_data, target_points, origin_points, interpolation_method)
182
+ elif len(source_data.shape) == 3:
183
+ interpolated_data = np.stack([interp_single(source_data[i], target_points, origin_points, interpolation_method) for i in range(source_data.shape[0])])
184
+ elif len(source_data.shape) == 4:
185
+ interpolated_data = np.stack([np.stack([interp_single(source_data[i, j], target_points, origin_points, interpolation_method) for j in range(source_data.shape[1])]) for i in range(source_data.shape[0])])
225
186
 
226
187
  return np.squeeze(np.array(interpolated_data))
227
188
 
228
189
 
229
- def mask_shapefile(data: np.ndarray, lons: np.ndarray, lats: np.ndarray, shapefile_path: str) -> xr.DataArray:
190
+ def mask_shapefile(
191
+ data_array: np.ndarray,
192
+ longitudes: np.ndarray,
193
+ latitudes: np.ndarray,
194
+ shapefile_path: str,
195
+ ) -> Union[xr.DataArray, None]:
230
196
  """
231
- Masks a 2D data array using a shapefile.
197
+ Mask a 2D data array using a shapefile.
232
198
 
233
- Parameters:
234
- - data: 2D numpy array of data to be masked.
235
- - lons: 1D numpy array of longitudes.
236
- - lats: 1D numpy array of latitudes.
237
- - shapefile_path: Path to the shapefile used for masking.
199
+ Args:
200
+ data_array (np.ndarray): 2D array of data to be masked.
201
+ longitudes (np.ndarray): 1D array of longitudes.
202
+ latitudes (np.ndarray): 1D array of latitudes.
203
+ shapefile_path (str): Path to the shapefile used for masking.
238
204
 
239
205
  Returns:
240
- - Masked xarray DataArray.
241
- """
242
- """
243
- https://cloud.tencent.com/developer/article/1701896
206
+ Union[xr.DataArray, None]: Masked xarray DataArray or None if an error occurs.
207
+
208
+ Raises:
209
+ FileNotFoundError: If the shapefile does not exist.
210
+ ValueError: If the data dimensions do not match the coordinates.
211
+
212
+ Examples:
213
+ >>> data_array = np.random.rand(10, 10)
214
+ >>> longitudes = np.linspace(-180, 180, 10)
215
+ >>> latitudes = np.linspace(-90, 90, 10)
216
+ >>> shapefile_path = "path/to/shapefile.shp"
217
+ >>> masked_data = mask_shapefile(data_array, longitudes, latitudes, shapefile_path)
218
+ >>> print(masked_data) # Expected output: Masked DataArray
219
+
244
220
  """
245
221
  try:
246
- # import geopandas as gpd
247
- # shp_f = gpd.read_file(shapefile_path)
248
222
  shp_f = salem.read_shapefile(shapefile_path)
249
- data_da = xr.DataArray(data, coords=[("latitude", lats), ("longitude", lons)])
223
+ data_da = xr.DataArray(data_array, coords=[("latitude", latitudes), ("longitude", longitudes)])
250
224
  masked_data = data_da.salem.roi(shape=shp_f)
251
225
  return masked_data
252
226
  except Exception as e:
253
- print(f"An error occurred: {e}")
227
+ print(f"[red]An error occurred: {e}[/red]")
254
228
  return None
255
229
 
256
230
 
257
231
  if __name__ == "__main__":
258
232
  pass
259
- """ import time
260
-
261
- import matplotlib.pyplot as plt
262
-
263
- # 测试数据
264
- origin_x = np.linspace(0, 10, 11)
265
- origin_y = np.linspace(0, 10, 11)
266
- target_x = np.linspace(0, 10, 101)
267
- target_y = np.linspace(0, 10, 101)
268
- data = np.random.rand(11, 11)
269
-
270
- # 高维插值
271
- origin_x = np.linspace(0, 10, 11)
272
- origin_y = np.linspace(0, 10, 11)
273
- target_x = np.linspace(0, 10, 101)
274
- target_y = np.linspace(0, 10, 101)
275
- data = np.random.rand(10, 10, 11, 11)
276
-
277
- start = time.time()
278
- interpolated_data = interp_2d(target_x, target_y, origin_x, origin_y, data, parallel=False)
279
- print(f"Interpolation time: {time.time()-start:.2f}s")
280
-
281
- print(interpolated_data.shape)
282
-
283
- # 高维插值多线程
284
- start = time.time()
285
- interpolated_data = interp_2d(target_x, target_y, origin_x, origin_y, data)
286
- print(f"Interpolation time: {time.time()-start:.2f}s")
287
-
288
- print(interpolated_data.shape)
289
- print(interpolated_data[0, 0, :, :].shape)
290
- plt.figure()
291
- plt.contourf(target_x, target_y, interpolated_data[0, 0, :, :])
292
- plt.colorbar()
293
- plt.show() """
oafuncs/oa_date.py CHANGED
@@ -17,62 +17,90 @@ import calendar
17
17
  import datetime
18
18
  from typing import List, Optional
19
19
 
20
+ from rich import print
21
+
20
22
  __all__ = ["month_days", "hour_range", "adjust_time", "timeit"]
21
23
 
22
24
 
23
25
  def month_days(year: int, month: int) -> int:
26
+ """
27
+ Calculate the number of days in a specific month of a year.
28
+
29
+ Args:
30
+ year (int): The year.
31
+ month (int): The month (1-12).
32
+
33
+ Returns:
34
+ int: Number of days in the specified month.
35
+
36
+ Example:
37
+ >>> month_days(2024, 2)
38
+ 29
39
+ """
24
40
  return calendar.monthrange(year, month)[1]
25
41
 
26
42
 
27
- def hour_range(start: str, end: str, interval: int = 6) -> List[str]:
43
+ def hour_range(start_time: str, end_time: str, hour_interval: int = 6) -> List[str]:
28
44
  """
29
45
  Generate a list of datetime strings with a specified interval in hours.
30
46
 
31
47
  Args:
32
- start (str): Start date in the format "%Y%m%d%H".
33
- end (str): End date in the format "%Y%m%d%H".
34
- interval (int): Interval in hours between each datetime.
48
+ start_time (str): Start date in the format "%Y%m%d%H".
49
+ end_time (str): End date in the format "%Y%m%d%H".
50
+ hour_interval (int): Interval in hours between each datetime.
35
51
 
36
52
  Returns:
37
- list: List of datetime strings in the format "%Y%m%d%H".
53
+ List[str]: List of datetime strings in the format "%Y%m%d%H".
54
+
55
+ Example:
56
+ >>> hour_range("2024010100", "2024010200", 6)
57
+ ['2024010100', '2024010106', '2024010112', '2024010118', '2024010124']
38
58
  """
39
- date_s = datetime.datetime.strptime(start, "%Y%m%d%H")
40
- date_e = datetime.datetime.strptime(end, "%Y%m%d%H")
59
+ date_s = datetime.datetime.strptime(start_time, "%Y%m%d%H")
60
+ date_e = datetime.datetime.strptime(end_time, "%Y%m%d%H")
41
61
  date_list = []
42
62
  while date_s <= date_e:
43
63
  date_list.append(date_s.strftime("%Y%m%d%H"))
44
- date_s += datetime.timedelta(hours=interval)
64
+ date_s += datetime.timedelta(hours=hour_interval)
45
65
  return date_list
46
66
 
47
67
 
48
- def adjust_time(initial_time: str, amount: int, time_unit: str = "hours", output_format: Optional[str] = None) -> str:
68
+ def adjust_time(base_time: str, time_delta: int, delta_unit: str = "hours", output_format: Optional[str] = None) -> str:
49
69
  """
50
- Adjust a given initial time by adding a specified amount of time.
70
+ Adjust a given base time by adding a specified time delta.
51
71
 
52
72
  Args:
53
- initial_time (str): Initial time in the format "yyyymmdd" to "yyyymmddHHMMSS".
54
- Missing parts are assumed to be "0".
55
- amount (int): The amount of time to add.
56
- time_unit (str): The unit of time to add ("seconds", "minutes", "hours", "days").
73
+ base_time (str): Base time in the format "yyyymmdd" to "yyyymmddHHMMSS".
74
+ Missing parts are assumed to be "0".
75
+ time_delta (int): The amount of time to add.
76
+ delta_unit (str): The unit of time to add ("seconds", "minutes", "hours", "days").
57
77
  output_format (str, optional): Custom output format for the adjusted time. Defaults to None.
58
78
 
59
79
  Returns:
60
80
  str: The adjusted time as a string, formatted according to the output_format or time unit.
81
+
82
+ Example:
83
+ >>> adjust_time("20240101", 5, "days")
84
+ '20240106'
85
+ >>> adjust_time("20240101000000", 2, "hours", "%Y-%m-%d %H:%M:%S")
86
+ '2024-01-01 02:00:00'
87
+ >>> adjust_time("20240101000000", 30, "minutes")
88
+ '2024-01-01 00:30:00'
61
89
  """
62
90
  # Normalize the input time to "yyyymmddHHMMSS" format
63
91
  time_format = "%Y%m%d%H%M%S"
64
- initial_time = initial_time.ljust(14, "0")
65
- time_obj = datetime.datetime.strptime(initial_time, time_format)
92
+ base_time = base_time.ljust(14, "0")
93
+ time_obj = datetime.datetime.strptime(base_time, time_format)
66
94
 
67
95
  # Add the specified amount of time
68
- if time_unit == "seconds":
69
- time_obj += datetime.timedelta(seconds=amount)
70
- elif time_unit == "minutes":
71
- time_obj += datetime.timedelta(minutes=amount)
72
- elif time_unit == "hours":
73
- time_obj += datetime.timedelta(hours=amount)
74
- elif time_unit == "days":
75
- time_obj += datetime.timedelta(days=amount)
96
+ if delta_unit == "seconds":
97
+ time_obj += datetime.timedelta(seconds=time_delta)
98
+ elif delta_unit == "minutes":
99
+ time_obj += datetime.timedelta(minutes=time_delta)
100
+ elif delta_unit == "hours":
101
+ time_obj += datetime.timedelta(hours=time_delta)
102
+ elif delta_unit == "days":
103
+ time_obj += datetime.timedelta(days=time_delta)
76
104
  else:
77
105
  raise ValueError("Invalid time unit. Use 'seconds', 'minutes', 'hours', or 'days'.")
78
106
 
@@ -80,13 +108,13 @@ def adjust_time(initial_time: str, amount: int, time_unit: str = "hours", output
80
108
  if output_format:
81
109
  return time_obj.strftime(output_format)
82
110
  else:
83
- if time_unit == "seconds":
111
+ if delta_unit == "seconds":
84
112
  default_format = "%Y%m%d%H%M%S"
85
- elif time_unit == "minutes":
113
+ elif delta_unit == "minutes":
86
114
  default_format = "%Y%m%d%H%M"
87
- elif time_unit == "hours":
115
+ elif delta_unit == "hours":
88
116
  default_format = "%Y%m%d%H"
89
- elif time_unit == "days":
117
+ elif delta_unit == "days":
90
118
  default_format = "%Y%m%d"
91
119
  return time_obj.strftime(default_format)
92
120
 
@@ -96,19 +124,25 @@ class timeit:
96
124
  A decorator to measure the execution time of a function.
97
125
 
98
126
  Usage:
99
- @timeit(log=True, print_time=True)
127
+ @timeit(log_to_file=True, display_time=True)
100
128
  def my_function():
101
129
  # Function code here
102
130
 
103
131
  Args:
104
- log (bool): Whether to log the execution time to a file. Defaults to False.
105
- print_time (bool): Whether to print the execution time to the console. Defaults to True.
132
+ log_to_file (bool): Whether to log the execution time to a file. Defaults to False.
133
+ display_time (bool): Whether to print the execution time to the console. Defaults to True.
134
+
135
+ Example:
136
+ @timeit(log_to_file=True, display_time=True)
137
+ def example_function():
138
+ # Simulate some work
139
+ time.sleep(2)
106
140
  """
107
141
 
108
- def __init__(self, func, log: bool = False, print_time: bool = True):
142
+ def __init__(self, func, log_to_file: bool = False, display_time: bool = True):
109
143
  self.func = func
110
- self.log = log
111
- self.print_time = print_time
144
+ self.log_to_file = log_to_file
145
+ self.display_time = display_time
112
146
 
113
147
  def __call__(self, *args, **kwargs):
114
148
  start_time = datetime.datetime.now()
@@ -116,10 +150,10 @@ class timeit:
116
150
  end_time = datetime.datetime.now()
117
151
  elapsed_time = (end_time - start_time).total_seconds()
118
152
 
119
- if self.print_time:
120
- print(f"Function '{self.func.__name__}' executed in {elapsed_time:.2f} seconds.")
153
+ if self.display_time:
154
+ print(f"[bold green]Function '{self.func.__name__}' executed in {elapsed_time:.2f} seconds.[/bold green]")
121
155
 
122
- if self.log:
156
+ if self.log_to_file:
123
157
  with open("execution_time.log", "a") as log_file:
124
158
  log_file.write(f"{datetime.datetime.now()} - Function '{self.func.__name__}' executed in {elapsed_time:.2f} seconds.\n")
125
159