amtypes.py

## ------------------------------------------
## Generated by Transwarp
##
## THIS FILE IS AUTOMATICALLY GENERATED.
## DO NOT EDIT. ALL CHANGES WILL BE LOST.
## ------------------------------------------

import struct
import enum
from dataclasses import dataclass, field
import socket

def _assert_equal(value, expected):
    if value != expected:
        raise ValueError(f"expected {expected} but found {value}")

class Zeroes:
    @classmethod
    def parse(cls, p, num):
        body = p[:num]
        _assert_equal(set(body) - {0}, set())
        return body, p[num:]

    def pack(self, num):
        return b"\x00" * num

class Bytes(bytes):
    @classmethod
    def parse(cls, p, size):
        return p[:size], p[size:]

    def pack(self, size):
        return self

class String(Bytes):
    @classmethod
    def parse(cls, p, size):
        body, p = super().parse(p, size)
        _assert_equal(body[-1], 0)
        return body[:-1].decode(), p

    def pack(self, size):
        return self[:size-1].ljust(size, '\x00').encode()

class Array:
    @classmethod
    def parse(cls, p, elem, num):
        res = []
        for _ in range(num):
            item, p = elem.parse(p)
            res.append(item)
        return res, p

    def pack(self, cls, num):
        return b"".join(cls.pack(e) for e in self)

class IPv4(str):
    @classmethod
    def parse(cls, p):
        addr = p[:4][::-1]
        return cls(socket.inet_ntoa(addr)), p[4:]

    def pack(self):
        return socket.inet_aton(self)[::-1]

class Magic(bytes):
    @classmethod
    def parse(cls, p, size, expected):
        v, p = p[:size], p[size:]
        _assert_equal(v, expected)
        return cls(v), p

    def pack(self, size, expected):
        return self

class Tail(bytes):
    @classmethod
    def parse(cls, p):
        return cls(p), b""

    def pack(self):
        if isinstance(self, bytes):
            return self
        else:
            return self.pack()

class IntType(int):
    pass

class i8be(IntType):
    size = 1

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack(">b", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack(">b", self)

class i8le(IntType):
    size = 1

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack("<b", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack("<b", self)

i8 = i8be

class u8be(IntType):
    size = 1

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack(">B", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack(">B", self)

class u8le(IntType):
    size = 1

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack("<B", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack("<B", self)

u8 = u8be

class i16be(IntType):
    size = 2

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack(">h", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack(">h", self)

class i16le(IntType):
    size = 2

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack("<h", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack("<h", self)

i16 = i16be

class u16be(IntType):
    size = 2

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack(">H", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack(">H", self)

class u16le(IntType):
    size = 2

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack("<H", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack("<H", self)

u16 = u16be

class i32be(IntType):
    size = 4

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack(">i", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack(">i", self)

class i32le(IntType):
    size = 4

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack("<i", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack("<i", self)

i32 = i32be

class u32be(IntType):
    size = 4

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack(">I", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack(">I", self)

class u32le(IntType):
    size = 4

    @classmethod
    def parse(cls, p):
        return cls(struct.unpack("<I", p[:cls.size])[0]), p[cls.size:]

    def pack(self):
        return struct.pack("<I", self)

u32 = u32be