Merge pull request #19 from MatrixEditor/fix/prefixed-encoding

[FIX] Prefixed `encoding` Parameter

examples/checksum_example.py

@@ -0,0 +1,23 @@
+from rich import print
+from cryptography.hazmat.primitives import hashes
+
+from caterpillar.py import struct, Bytes, pack, unpack, this
+from caterpillar.fields.digest import HMAC
+
+
+@struct
+class Format:
+    key: b"MAGIC"
+
+    with HMAC(this.key, hashes.SHA256(), "hmac", verify=True):
+        user_data: Bytes(11)
+
+    # attribute 'hmac' is added by HMAC
+
+
+print(Format(user_data=b"hello world"))
+print(Format.__struct__.fields)
+data = pack(Format(user_data=b"hello world"))
+print(data)
+print("unpacking..")
+print(unpack(Format, data[:-1] + b"\x00"))

examples/formats/nibarchive.py

@@ -87,7 +87,7 @@ class NIBClassName:
 
 # Note that the returned string instance here may contain extra null-bytes
 # at the end.
-NIBKey = Prefixed(vint, String(...))
+NIBKey = Prefixed(vint, encoding="utf-8")
 
 
 class ValueType(enum.Enum):

src/caterpillar/_common.py

@@ -14,14 +14,13 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 import itertools
 
-from typing import List, Any, Union
+from typing import Collection, List, Any
 
-from caterpillar.abc import _GreedyType, _ContextLike, _StreamType, _PrefixedType
+from caterpillar.abc import _ContextLike, _StreamType, _PrefixedType
 from caterpillar.context import (
     Context,
     CTX_PATH,
     CTX_FIELD,
-    CTX_POS,
     CTX_INDEX,
     CTX_OBJECT,
     CTX_STREAM,
@@ -63,7 +62,7 @@ def unpack_seq(context: _ContextLike, unpack_one) -> List[Any]:
     field = context[CTX_FIELD]
     assert field and context[CTX_SEQ]
 
-    length: Union[int, _GreedyType] = field.length(context)
+    length = field.length(context)
     base_path = context[CTX_PATH]
     # Special elements '_index' and '_length' can be referenced within
     # the new context. The '_pos' attribute will be adjusted automatically.
@@ -114,7 +113,7 @@ def unpack_seq(context: _ContextLike, unpack_one) -> List[Any]:
     return values
 
 
-def pack_seq(seq: List[Any], context: _ContextLike, pack_one) -> None:
+def pack_seq(seq: Collection[Any], context: _ContextLike, pack_one) -> None:
     """Generic function to pack sequenced elements.
 
     :param seq: the iterable of elements

src/caterpillar/fields/__init__.py

@@ -83,5 +83,5 @@
     Sha3_512,
     Crc32,
     Adler,
-    HMAC
+    HMAC,
 )

src/caterpillar/fields/_mixin.py

@@ -12,15 +12,18 @@
 #
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
-import inspect
-
 from io import BytesIO
-from typing import Any, List, Union, Optional, Iterable, Callable
-from abc import ABC, abstractmethod
+from types import EllipsisType
+from typing import Any, Collection, List, Union, Iterable, Callable
 from functools import partial
 
-from caterpillar.abc import _ContextLike, _StructLike, _ContextLambda, _Switch
-from caterpillar.abc import _ContainsStruct, getstruct
+from caterpillar.abc import (
+    _ContextLike,
+    _StructLike,
+    _ContextLambda,
+    _Switch,
+    getstruct,
+)
 from caterpillar.byteorder import ByteOrder, byteorder
 from caterpillar.options import Flag
 from caterpillar.context import CTX_SEQ, CTX_STREAM
@@ -44,7 +47,7 @@ def __matmul__(self, offset: Union[_ContextLambda, int]) -> Field:
         """Creates a field that should start at the given offset."""
         return Field(self, byteorder(self)) @ offset
 
-    def __getitem__(self, dim: Union[_ContextLambda, int]) -> Field:
+    def __getitem__(self, dim: Union[_ContextLambda, int, EllipsisType]) -> Field:
         """Returns a sequenced field."""
         return Field(self, byteorder(self))[dim]
 
@@ -113,7 +116,7 @@ def unpack_single(self, context: _ContextLike) -> Any:
         """
         raise NotImplementedError
 
-    def pack_seq(self, seq: Iterable, context: _ContextLike) -> None:
+    def pack_seq(self, seq: Collection, context: _ContextLike) -> None:
         """
         Pack a sequence of elements using the provided context.
 

src/caterpillar/fields/common.py

@@ -1141,7 +1141,7 @@ def __init__(
     ):
         self.prefix = prefix
         self.struct = struct
-        self.encoding = None
+        self.encoding = encoding
 
     def __type__(self) -> Optional[Union[type, str]]:
         """

src/caterpillar/fields/crypto.py

@@ -18,7 +18,7 @@
     from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
     from cryptography.hazmat.primitives.padding import PaddingContext
 except ImportError:
-    Cipher = algorithms = modes = PaddingContext = None
+    Cipher = algorithms = modes = PaddingContext = Any
 
 
 from caterpillar.abc import _StructLike, _ContextLike

test/_Py/fields/test_py_prefixed.py

@@ -0,0 +1,29 @@
+from caterpillar.py import Prefixed, pack, unpack, uint8, Bytes, uint16
+
+
+def test_prefixed_bytes():
+    atom = Prefixed(uint8)
+
+    assert pack(b"foo", atom, as_field=True) == b"\x03foo"
+    assert unpack(atom, b"\x03foo", as_field=True) == b"foo"
+
+
+def test_prefixed_encoding():
+    atom = Prefixed(uint8, encoding="utf-8")
+
+    assert unpack(atom, b"\x03foo", as_field=True) == "foo"
+    assert pack("foo", atom, as_field=True) == b"\x03foo"
+
+
+def test_prefixed_bytes_struct():
+    atom = Prefixed(uint8, Bytes(...))
+
+    assert unpack(atom, b"\x03foo", as_field=True) == b"foo"
+    assert pack(b"foo", atom, as_field=True) == b"\x03foo"
+
+
+def test_prefixed_custom_struct():
+    atom = Prefixed(uint8, uint16[...])
+
+    assert unpack(atom, b"\x04\x01\x00\x01\x00", as_field=True) == [1, 1]
+    assert pack([1, 1], atom, as_field=True) == b"\x04\x01\x00\x01\x00"