def read_command()

in antlir/btrfs_diff/parse_send_stream.py [0:0]

• 118 lines of code
• 27 McCabe index (conditional complexity)

def read_command(infile):
    cmd_header = CommandHeader.from_file(infile)

    s = infile.read(cmd_header.length)
    if len(s) != cmd_header.length:
        raise RuntimeError(f"{cmd_header} got {len(s)} bytes")
    # Future: pull in the `crc32c` module and check the CRC.

    attr_bytes = BytesIO(s)
    kind_to_attr = {}
    while attr_bytes.tell() != len(s):
        kind, attr = read_attribute(attr_bytes)
        if kind in kind_to_attr:
            raise RuntimeError(f"{kind} occurred twice in {cmd_header}")
        kind_to_attr[kind] = attr

    if cmd_header.kind == CommandKind.SUBVOL:
        return SendStreamItems.subvol(
            path=kind_to_attr[AttributeKind.PATH],
            uuid=kind_to_attr[AttributeKind.UUID],
            transid=kind_to_attr[AttributeKind.CTRANSID],
        )
    elif cmd_header.kind == CommandKind.SNAPSHOT:
        return SendStreamItems.snapshot(
            path=kind_to_attr[AttributeKind.PATH],
            uuid=kind_to_attr[AttributeKind.UUID],
            transid=kind_to_attr[AttributeKind.CTRANSID],
            parent_uuid=kind_to_attr[AttributeKind.CLONE_UUID],
            parent_transid=kind_to_attr[AttributeKind.CLONE_CTRANSID],
        )
    elif cmd_header.kind == CommandKind.MKFILE:
        return SendStreamItems.mkfile(path=kind_to_attr[AttributeKind.PATH])
    elif cmd_header.kind == CommandKind.MKDIR:
        return SendStreamItems.mkdir(path=kind_to_attr[AttributeKind.PATH])
    elif cmd_header.kind == CommandKind.MKNOD:
        return SendStreamItems.mknod(
            path=kind_to_attr[AttributeKind.PATH],
            mode=kind_to_attr[AttributeKind.MODE],
            dev=kind_to_attr[AttributeKind.RDEV],
        )
    elif cmd_header.kind == CommandKind.MKFIFO:
        return SendStreamItems.mkfifo(path=kind_to_attr[AttributeKind.PATH])
    elif cmd_header.kind == CommandKind.MKSOCK:
        return SendStreamItems.mksock(path=kind_to_attr[AttributeKind.PATH])
    elif cmd_header.kind == CommandKind.SYMLINK:
        return SendStreamItems.symlink(
            path=kind_to_attr[AttributeKind.PATH],
            # NB Unlike the other `dest` attributes, we don't normalize this.
            dest=os.path.normpath(kind_to_attr[AttributeKind.PATH_LINK]),
        )
    elif cmd_header.kind == CommandKind.RENAME:
        return SendStreamItems.rename(
            path=kind_to_attr[AttributeKind.PATH],
            dest=kind_to_attr[AttributeKind.PATH_TO],
        )
    elif cmd_header.kind == CommandKind.LINK:
        return SendStreamItems.link(
            path=kind_to_attr[AttributeKind.PATH],
            dest=os.path.normpath(kind_to_attr[AttributeKind.PATH_LINK]),
        )
    elif cmd_header.kind == CommandKind.UNLINK:
        return SendStreamItems.unlink(path=kind_to_attr[AttributeKind.PATH])
    elif cmd_header.kind == CommandKind.RMDIR:
        return SendStreamItems.rmdir(path=kind_to_attr[AttributeKind.PATH])
    elif cmd_header.kind == CommandKind.WRITE:
        return SendStreamItems.write(
            path=kind_to_attr[AttributeKind.PATH],
            offset=kind_to_attr[AttributeKind.FILE_OFFSET],
            data=kind_to_attr[AttributeKind.DATA],
        )
    elif cmd_header.kind == CommandKind.CLONE:
        return SendStreamItems.clone(
            path=kind_to_attr[AttributeKind.PATH],
            offset=kind_to_attr[AttributeKind.FILE_OFFSET],
            len=kind_to_attr[AttributeKind.CLONE_LEN],
            from_uuid=kind_to_attr[AttributeKind.CLONE_UUID],
            from_transid=kind_to_attr[AttributeKind.CLONE_CTRANSID],
            from_path=kind_to_attr[AttributeKind.CLONE_PATH],
            clone_offset=kind_to_attr[AttributeKind.CLONE_OFFSET],
        )
    elif cmd_header.kind == CommandKind.SET_XATTR:
        return SendStreamItems.set_xattr(
            path=kind_to_attr[AttributeKind.PATH],
            name=kind_to_attr[AttributeKind.XATTR_NAME],
            data=kind_to_attr[AttributeKind.XATTR_DATA],
        )
    elif cmd_header.kind == CommandKind.REMOVE_XATTR:
        return SendStreamItems.remove_xattr(
            path=kind_to_attr[AttributeKind.PATH],
            name=kind_to_attr[AttributeKind.XATTR_NAME],
        )
    elif cmd_header.kind == CommandKind.TRUNCATE:
        return SendStreamItems.truncate(
            path=kind_to_attr[AttributeKind.PATH],
            size=kind_to_attr[AttributeKind.SIZE],
        )
    elif cmd_header.kind == CommandKind.CHMOD:
        return SendStreamItems.chmod(
            path=kind_to_attr[AttributeKind.PATH],
            mode=kind_to_attr[AttributeKind.MODE],
        )
    elif cmd_header.kind == CommandKind.CHOWN:
        return SendStreamItems.chown(
            path=kind_to_attr[AttributeKind.PATH],
            uid=kind_to_attr[AttributeKind.UID],
            gid=kind_to_attr[AttributeKind.GID],
        )
    elif cmd_header.kind == CommandKind.UTIMES:
        return SendStreamItems.utimes(
            path=kind_to_attr[AttributeKind.PATH],
            ctime=kind_to_attr[AttributeKind.CTIME],
            mtime=kind_to_attr[AttributeKind.MTIME],
            atime=kind_to_attr[AttributeKind.ATIME],
        )
    elif cmd_header.kind == CommandKind.END:
        return None
    elif cmd_header.kind == CommandKind.UPDATE_EXTENT:
        return SendStreamItems.update_extent(
            path=kind_to_attr[AttributeKind.PATH],
            offset=kind_to_attr[AttributeKind.FILE_OFFSET],
            len=kind_to_attr[AttributeKind.SIZE],
        )

    raise AssertionError(f"Fix me: unhandled {cmd_header}")  # pragma: no cover