Skip to content

Commit 3ef7b26

Browse files
caffeinismcaffeinism
committed
ruff
1 parent 56dd2dc commit 3ef7b26

File tree

2 files changed

+80
-28
lines changed

2 files changed

+80
-28
lines changed

‎av/video/frame.py‎

Lines changed: 39 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
_cuda_device_ctx_cache = {}
2929
_cuda_frames_ctx_cache = {}
3030

31+
3132
@cython.cfunc
3233
def _consume_dlpack(obj: object, stream: object) -> cython.pointer[DLManagedTensor]:
3334
capsule: object
@@ -40,7 +41,9 @@ def _consume_dlpack(obj: object, stream: object) -> cython.pointer[DLManagedTens
4041

4142
if not PyCapsule_IsValid(capsule, b"dltensor"):
4243
PyErr_Clear()
43-
raise TypeError("expected a DLPack capsule or an object implementing __dlpack__")
44+
raise TypeError(
45+
"expected a DLPack capsule or an object implementing __dlpack__"
46+
)
4447

4548
managed = cython.cast(
4649
cython.pointer[DLManagedTensor],
@@ -54,6 +57,7 @@ def _consume_dlpack(obj: object, stream: object) -> cython.pointer[DLManagedTens
5457

5558
return managed
5659

60+
5761
@cython.cfunc
5862
@cython.nogil
5963
@cython.exceptval(check=False)
@@ -67,6 +71,7 @@ def _dlpack_avbuffer_free(
6771
if managed != cython.NULL:
6872
managed.deleter(managed)
6973

74+
7075
@cython.cfunc
7176
def _get_cuda_device_ctx(
7277
device_id: cython.int,
@@ -103,6 +108,7 @@ def _get_cuda_device_ctx(
103108
_cuda_device_ctx_cache[key] = cython.cast(cython.size_t, device_ref)
104109
return device_ref
105110

111+
106112
@cython.cfunc
107113
def _get_cuda_frames_ctx(
108114
device_id: cython.int,
@@ -398,7 +404,9 @@ def planes(self):
398404
frames_ctx: cython.pointer[AVHWFramesContext] = cython.cast(
399405
cython.pointer[AVHWFramesContext], self.ptr.hw_frames_ctx.data
400406
)
401-
fmt = get_video_format(frames_ctx.sw_format, self.ptr.width, self.ptr.height)
407+
fmt = get_video_format(
408+
frames_ctx.sw_format, self.ptr.width, self.ptr.height
409+
)
402410

403411
max_plane_count: cython.int = 0
404412
for i in range(fmt.ptr.nb_components):
@@ -1347,7 +1355,9 @@ def from_dlpack(
13471355
planes = (planes,)
13481356

13491357
if len(planes) != 2:
1350-
raise ValueError("from_dlpack currently supports 2-plane formats only (nv12/p010le/p016le)")
1358+
raise ValueError(
1359+
"from_dlpack currently supports 2-plane formats only (nv12/p010le/p016le)"
1360+
)
13511361

13521362
sw_fmt: lib.AVPixelFormat = get_pix_fmt(format)
13531363
nv12 = get_pix_fmt(b"nv12")
@@ -1373,7 +1383,9 @@ def from_dlpack(
13731383
if dev_type0 != dev_type1:
13741384
raise ValueError("plane tensors must have the same device_type")
13751385
if dev_type0 not in {kDLCUDA, kDLCPU}:
1376-
raise NotImplementedError("only CPU and CUDA DLPack tensors are supported")
1386+
raise NotImplementedError(
1387+
"only CPU and CUDA DLPack tensors are supported"
1388+
)
13771389

13781390
dev0 = m0.dl_tensor.device.device_id
13791391
dev1 = m1.dl_tensor.device.device_id
@@ -1385,7 +1397,9 @@ def from_dlpack(
13851397
if device_id is None:
13861398
device_id = dev0
13871399
elif device_id != dev0:
1388-
raise ValueError("device_id does not match the DLPack tensor device_id")
1400+
raise ValueError(
1401+
"device_id does not match the DLPack tensor device_id"
1402+
)
13891403
else:
13901404
if device_id not in (None, 0):
13911405
raise ValueError("device_id must be 0 for CPU tensors")
@@ -1445,7 +1459,9 @@ def from_dlpack(
14451459
raise ValueError("plane 1 must have shape (H/2, W) for 2D UV")
14461460
if m1.dl_tensor.strides != cython.NULL:
14471461
if m1.dl_tensor.strides[1] != 1:
1448-
raise ValueError("plane 1 must be contiguous in the last dimension")
1462+
raise ValueError(
1463+
"plane 1 must be contiguous in the last dimension"
1464+
)
14491465
uv_pitch_elems = cython.cast(int64_t, m1.dl_tensor.strides[0])
14501466
else:
14511467
uv_pitch_elems = cython.cast(int64_t, uv_w)
@@ -1457,7 +1473,9 @@ def from_dlpack(
14571473
raise ValueError("plane 1 must have shape (H/2, W/2, 2) for 3D UV")
14581474
if m1.dl_tensor.strides != cython.NULL:
14591475
if m1.dl_tensor.strides[2] != 1 or m1.dl_tensor.strides[1] != 2:
1460-
raise ValueError("unexpected UV plane strides for (H/2, W/2, 2)")
1476+
raise ValueError(
1477+
"unexpected UV plane strides for (H/2, W/2, 2)"
1478+
)
14611479
uv_pitch_elems = cython.cast(int64_t, m1.dl_tensor.strides[0])
14621480
else:
14631481
uv_pitch_elems = cython.cast(int64_t, width)
@@ -1477,7 +1495,9 @@ def from_dlpack(
14771495
raise TypeError("primary_ctx must be a bool")
14781496
primary_ctx = bool(primary_ctx)
14791497

1480-
frames_ref = _get_cuda_frames_ctx(device_id, primary_ctx, sw_fmt, width, height)
1498+
frames_ref = _get_cuda_frames_ctx(
1499+
device_id, primary_ctx, sw_fmt, width, height
1500+
)
14811501

14821502
frame.ptr.format = get_pix_fmt(b"cuda")
14831503
frame.ptr.hw_frames_ctx = lib.av_buffer_ref(frames_ref)
@@ -1486,12 +1506,12 @@ def from_dlpack(
14861506
else:
14871507
frame.ptr.format = sw_fmt
14881508

1489-
y_ptr = cython.cast(cython.pointer[uint8_t], m0.dl_tensor.data) + cython.cast(
1490-
cython.size_t, m0.dl_tensor.byte_offset
1491-
)
1492-
uv_ptr = cython.cast(cython.pointer[uint8_t], m1.dl_tensor.data) + cython.cast(
1493-
cython.size_t, m1.dl_tensor.byte_offset
1494-
)
1509+
y_ptr = cython.cast(
1510+
cython.pointer[uint8_t], m0.dl_tensor.data
1511+
) + cython.cast(cython.size_t, m0.dl_tensor.byte_offset)
1512+
uv_ptr = cython.cast(
1513+
cython.pointer[uint8_t], m1.dl_tensor.data
1514+
) + cython.cast(cython.size_t, m1.dl_tensor.byte_offset)
14951515

14961516
frame.ptr.buf[0] = lib.av_buffer_create(
14971517
y_ptr, y_size, _dlpack_avbuffer_free, cython.cast(cython.p_void, m0), 0
@@ -1503,7 +1523,11 @@ def from_dlpack(
15031523
m0 = cython.NULL
15041524

15051525
frame.ptr.buf[1] = lib.av_buffer_create(
1506-
uv_ptr, uv_size, _dlpack_avbuffer_free, cython.cast(cython.p_void, m1), 0
1526+
uv_ptr,
1527+
uv_size,
1528+
_dlpack_avbuffer_free,
1529+
cython.cast(cython.p_void, m1),
1530+
0,
15071531
)
15081532
if frame.ptr.buf[1] == cython.NULL:
15091533
raise MemoryError("av_buffer_create failed for plane 1")

‎av/video/plane.py‎

Lines changed: 41 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,9 @@ def __cinit__(self, frame: VideoFrame, index: cython.int):
2929
frames_ctx: cython.pointer[AVHWFramesContext] = cython.cast(
3030
cython.pointer[AVHWFramesContext], frame.ptr.hw_frames_ctx.data
3131
)
32-
fmt = get_video_format(frames_ctx.sw_format, frame.ptr.width, frame.ptr.height)
32+
fmt = get_video_format(
33+
frames_ctx.sw_format, frame.ptr.width, frame.ptr.height
34+
)
3335

3436
if fmt.name == "pal8" and index == 1:
3537
self.width = 256
@@ -82,8 +84,12 @@ def __getbuffer__(self, view: cython.pointer[Py_buffer], flags: cython.int):
8284

8385
def __dlpack_device__(self):
8486
if self.frame.ptr.hw_frames_ctx:
85-
if cython.cast(lib.AVPixelFormat, self.frame.ptr.format) != get_pix_fmt(b"cuda"):
86-
raise NotImplementedError("DLPack export is only implemented for CUDA hw frames")
87+
if cython.cast(lib.AVPixelFormat, self.frame.ptr.format) != get_pix_fmt(
88+
b"cuda"
89+
):
90+
raise NotImplementedError(
91+
"DLPack export is only implemented for CUDA hw frames"
92+
)
8793

8894
frames_ctx: cython.pointer[AVHWFramesContext] = cython.cast(
8995
cython.pointer[AVHWFramesContext], self.frame.ptr.hw_frames_ctx.data
@@ -97,15 +103,21 @@ def __dlpack_device__(self):
97103

98104
def __dlpack__(self, stream=None):
99105
if self.frame.ptr.buf[0] == cython.NULL:
100-
raise TypeError("DLPack export requires a refcounted AVFrame (frame.buf[0] is NULL)")
106+
raise TypeError(
107+
"DLPack export requires a refcounted AVFrame (frame.buf[0] is NULL)"
108+
)
101109

102110
device_type: cython.int
103111
device_id: cython.int
104112
sw_fmt: lib.AVPixelFormat
105113

106114
if self.frame.ptr.hw_frames_ctx:
107-
if cython.cast(lib.AVPixelFormat, self.frame.ptr.format) != get_pix_fmt(b"cuda"):
108-
raise NotImplementedError("DLPack export is only implemented for CUDA hw frames")
115+
if cython.cast(lib.AVPixelFormat, self.frame.ptr.format) != get_pix_fmt(
116+
b"cuda"
117+
):
118+
raise NotImplementedError(
119+
"DLPack export is only implemented for CUDA hw frames"
120+
)
109121

110122
frames_ctx: cython.pointer[AVHWFramesContext] = cython.cast(
111123
cython.pointer[AVHWFramesContext], self.frame.ptr.hw_frames_ctx.data
@@ -122,7 +134,9 @@ def __dlpack__(self, stream=None):
122134

123135
line_size = self.line_size
124136
if line_size < 0:
125-
raise NotImplementedError("negative linesize is not supported for DLPack export")
137+
raise NotImplementedError(
138+
"negative linesize is not supported for DLPack export"
139+
)
126140

127141
nv12 = get_pix_fmt(b"nv12")
128142
p010le = get_pix_fmt(b"p010le")
@@ -187,8 +201,12 @@ def __dlpack__(self, stream=None):
187201
raise MemoryError("av_frame_alloc() failed")
188202
err_check(lib.av_frame_ref(frame_ref, self.frame.ptr))
189203

190-
shape = cython.cast(cython.pointer[int64_t], malloc(ndim * cython.sizeof(int64_t)))
191-
strides = cython.cast(cython.pointer[int64_t], malloc(ndim * cython.sizeof(int64_t)))
204+
shape = cython.cast(
205+
cython.pointer[int64_t], malloc(ndim * cython.sizeof(int64_t))
206+
)
207+
strides = cython.cast(
208+
cython.pointer[int64_t], malloc(ndim * cython.sizeof(int64_t))
209+
)
192210
if shape == cython.NULL or strides == cython.NULL:
193211
if shape != cython.NULL:
194212
free(shape)
@@ -210,7 +228,9 @@ def __dlpack__(self, stream=None):
210228
strides[1] = st1
211229
strides[2] = st2
212230

213-
ctx = cython.cast(cython.pointer[cython.p_void], malloc(3 * cython.sizeof(cython.p_void)))
231+
ctx = cython.cast(
232+
cython.pointer[cython.p_void], malloc(3 * cython.sizeof(cython.p_void))
233+
)
214234
if ctx == cython.NULL:
215235
free(shape)
216236
free(strides)
@@ -221,7 +241,9 @@ def __dlpack__(self, stream=None):
221241
ctx[1] = cython.cast(cython.p_void, shape)
222242
ctx[2] = cython.cast(cython.p_void, strides)
223243

224-
managed = cython.cast(cython.pointer[DLManagedTensor], malloc(cython.sizeof(DLManagedTensor)))
244+
managed = cython.cast(
245+
cython.pointer[DLManagedTensor], malloc(cython.sizeof(DLManagedTensor))
246+
)
225247
if managed == cython.NULL:
226248
free(ctx)
227249
free(shape)
@@ -243,7 +265,11 @@ def __dlpack__(self, stream=None):
243265
managed.deleter = _dlpack_managed_tensor_deleter
244266

245267
try:
246-
capsule = PyCapsule_New(cython.cast(cython.p_void, managed), b"dltensor", _dlpack_capsule_destructor)
268+
capsule = PyCapsule_New(
269+
cython.cast(cython.p_void, managed),
270+
b"dltensor",
271+
_dlpack_capsule_destructor,
272+
)
247273
except Exception:
248274
_dlpack_managed_tensor_deleter(managed)
249275
raise
@@ -254,7 +280,9 @@ def __dlpack__(self, stream=None):
254280
@cython.cfunc
255281
@cython.nogil
256282
@cython.exceptval(check=False)
257-
def _dlpack_managed_tensor_deleter(managed: cython.pointer[DLManagedTensor]) -> cython.void:
283+
def _dlpack_managed_tensor_deleter(
284+
managed: cython.pointer[DLManagedTensor],
285+
) -> cython.void:
258286
ctx: cython.pointer[cython.p_void]
259287
frame_ref: cython.pointer[lib.AVFrame]
260288
shape: cython.pointer[int64_t]

0 commit comments

Comments
 (0)