1try:
2    from collections.abc import Sequence
3except ImportError:
4    from collections import Sequence
5
6
7cdef object _cinit_bypass_sentinel = object()
8
9
10# Cython doesn't let us inherit from the abstract Sequence, so we will subclass
11# it later.
12cdef class _MotionVectors(SideData):
13
14    def __init__(self, *args, **kwargs):
15        super().__init__(*args, **kwargs)
16        self._vectors = {}
17        self._len = self.ptr.size // sizeof(lib.AVMotionVector)
18
19    def __repr__(self):
20        return f'<av.sidedata.MotionVectors {self.ptr.size} bytes of {len(self)} vectors at 0x{<unsigned int>self.ptr.data:0x}'
21
22    def __getitem__(self, int index):
23
24        try:
25            return self._vectors[index]
26        except KeyError:
27            pass
28
29        if index >= self._len:
30            raise IndexError(index)
31
32        vector = self._vectors[index] = MotionVector(_cinit_bypass_sentinel, self, index)
33        return vector
34
35    def __len__(self):
36        return self._len
37
38    def to_ndarray(self):
39        import numpy as np
40        return np.frombuffer(self, dtype=np.dtype([
41            ('source', 'int32'),
42            ('w', 'uint8'),
43            ('h', 'uint8'),
44            ('src_x', 'int16'),
45            ('src_y', 'int16'),
46            ('dst_x', 'int16'),
47            ('dst_y', 'int16'),
48            ('flags', 'uint64'),
49            ('motion_x', 'int32'),
50            ('motion_y', 'int32'),
51            ('motion_scale', 'uint16'),
52        ], align=True))
53
54
55class MotionVectors(_MotionVectors, Sequence):
56    pass
57
58
59cdef class MotionVector(object):
60
61    def __init__(self, sentinel, _MotionVectors parent, int index):
62        if sentinel is not _cinit_bypass_sentinel:
63            raise RuntimeError('cannot manually instatiate MotionVector')
64        self.parent = parent
65        cdef lib.AVMotionVector *base = <lib.AVMotionVector*>parent.ptr.data
66        self.ptr = base + index
67
68    def __repr__(self):
69        return f'<av.sidedata.MotionVector {self.w}x{self.h} from {self.src_x},{self.src_y} to {self.dst_x},{self.dst_y}>'
70
71    @property
72    def source(self):
73        return self.ptr.source
74
75    @property
76    def w(self):
77        return self.ptr.w
78
79    @property
80    def h(self):
81        return self.ptr.h
82
83    @property
84    def src_x(self):
85        return self.ptr.src_x
86
87    @property
88    def src_y(self):
89        return self.ptr.src_y
90
91    @property
92    def dst_x(self):
93        return self.ptr.dst_x
94
95    @property
96    def dst_y(self):
97        return self.ptr.dst_y
98
99    @property
100    def motion_x(self):
101        return self.ptr.motion_x
102
103    @property
104    def motion_y(self):
105        return self.ptr.motion_y
106
107    @property
108    def motion_scale(self):
109        return self.ptr.motion_scale
110