xref: /aosp_15_r20/external/executorch/exir/_serialize/_cord.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7import io
8from typing import List, Optional, Union
9
10
11class Cord:
12    """A `bytes`-like sequence of bytes, stored non-contiguously.
13
14    Users can use a Cord to assemble large files and data blobs using references
15    to and slices of other data, instead of copying and appending that data to a
16    `bytes` or `bytearray` object.
17    """
18
19    def __init__(self, data: Optional[Union[bytes, "Cord"]] = None) -> None:
20        """Initialize Cord data structure."""
21        self._buffers: List[bytes] = []
22        self._byte_size: int = 0
23
24        if data is not None:
25            self.append(data)
26
27    def __len__(self):
28        """Number of bytes in the Cord."""
29        return self._byte_size
30
31    def __bytes__(self) -> bytes:
32        """Return the contents of the Cord as a single `bytes` object."""
33        return b"".join(self._buffers)
34
35    def append(self, data: Union[bytes, "Cord"]) -> None:
36        """Append a bytes or Cord to the current Cord."""
37        if isinstance(data, bytes):
38            self._buffers.append(data)
39            self._byte_size += len(data)
40        elif isinstance(data, Cord):
41            self._buffers.extend(data._buffers)
42            self._byte_size += len(data)
43        else:
44            raise TypeError(f"Can only append bytes or Cords, received {type(data)}")
45
46    def write_to_file(self, outfile: io.BufferedIOBase) -> None:
47        """Write the Cord to a file."""
48        for item in self._buffers:
49            outfile.write(item)
50