def _fill_na(cls, bars: np.array, n: int, end: Frame, frame_type) -> np.ndarray:
if frame_type in tf.minute_level_frames:
convert = tf.int2time
else:
convert = tf.int2date
frames = [convert(x) for x in tf.get_frames_by_count(end, n, frame_type)]
filled = np.empty(n, dtype=bars.dtype)
filled[:] = np.nan
filled["frame"] = frames
return merge(filled, bars, "frame")
async def get_bars(
self,
code: str,
end: Union[datetime.date, datetime.datetime, Arrow],
n: int,
frame_type: FrameType,
) -> np.ndarray:
if n == 0:
return np.ndarray([], dtype=bars_dtype)
frames = tf.get_frames_by_count(end, n, frame_type)
tr = self.security.pipeline()
key = f"{code}:{frame_type.value}"
[tr.hget(key, int(frame)) for frame in frames]
recs = await tr.execute()
converter = (tf.int2time if frame_type in [
FrameType.MIN1,
FrameType.MIN5,
FrameType.MIN15,
FrameType.MIN30,
FrameType.MIN60,
] else tf.int2date)
data = np.empty(len(frames), dtype=bars_dtype)
for i, frame in enumerate(frames):
rec = recs[i]
if rec is None:
data[i] = (converter(frame), None, None, None, None, None,
None, None)
else:
o, h, l, c, v, a, f = rec.split(" ")
data[i] = (
converter(frame),
float(o),
float(h),
float(l),
float(c),
float(v),
float(a),
float(f),
)
return data
async def get_bars_raw_data(
self,
code: str,
end: Union[datetime.date, datetime.datetime, Arrow],
n: int,
frame_type: FrameType,
) -> bytes:
"""
如果没有数据，返回空字节串''
"""
if n == 0:
return b""
frames = tf.get_frames_by_count(end, n, frame_type)
pl = self.security.pipeline()
key = f"{code}:{frame_type.value}"
[pl.hget(key, int(frame), encoding=None) for frame in frames]
recs = await pl.execute()
return b"".join(filter(None, recs))
def test_get_frames_by_count(self):
days = [
20200117,
20200120,
20200121,
20200122,
20200123,
20200203,
20200204,
20200205,
20200206,
20200207,
20200210,
20200211,
]
for i in range(len(days)):
end, n = tf.int2date(days[i]), i + 1
expected = days[:n]
actual = tf.get_frames_by_count(end, n, FrameType.DAY)
logger.debug(
"get_frames_by_count(%s, %s, %s)->%s", end, n, FrameType.DAY, actual
)
self.assertListEqual(expected, list(actual))
X = [
(202002041030, 1, [202002041030]),
(202002041030, 2, [202002041000, 202002041030]),
(202002041030, 3, [202002031500, 202002041000, 202002041030]),
(202002041030, 4, [202002031430, 202002031500, 202002041000, 202002041030]),
(
202002041030,
5,
[202002031400, 202002031430, 202002031500, 202002041000, 202002041030],
),
(
202002041030,
6,
[
202002031330,
202002031400,
202002031430,
202002031500,
202002041000,
202002041030,
],
),
(
202002041030,
7,
[
202002031130,
202002031330,
202002031400,
202002031430,
202002031500,
202002041000,
202002041030,
],
),
(
202002041030,
8,
[
202002031100,
202002031130,
202002031330,
202002031400,
202002031430,
202002031500,
202002041000,
202002041030,
],
),
(
202002041030,
9,
[
202002031030,
202002031100,
202002031130,
202002031330,
202002031400,
202002031430,
202002031500,
202002041000,
202002041030,
],
),
(
202002041030,
10,
[
202002031000,
202002031030,
202002031100,
202002031130,
202002031330,
202002031400,
202002031430,
202002031500,
202002041000,
202002041030,
],
),
(
202002041030,
11,
[
202001231500,
202002031000,
202002031030,
202002031100,
202002031130,
202002031330,
202002031400,
202002031430,
202002031500,
202002041000,
202002041030,
],
),
]
for i, (end, n, expected) in enumerate(X):
end = tf.int2time(end)
actual = tf.get_frames_by_count(end, n, FrameType.MIN30)
logger.debug(
"get_frames_by_count(%s, %s, %s)->%s", end, n, FrameType.DAY, actual
)
self.assertListEqual(expected, actual)
actual = tf.get_frames_by_count(datetime.date(2020, 2, 12), 3, FrameType.MONTH)
self.assertListEqual([20191129, 20191231, 20200123], actual.tolist())
actual = tf.get_frames_by_count(datetime.date(2020, 2, 12), 3, FrameType.WEEK)
self.assertListEqual([20200117, 20200123, 20200207], actual.tolist())