tartrazine/python/enry/utils.py

from _c_enry import ffi
from enry.types import Guess
from functools import wraps
from typing import Hashable, List, Sequence


def py_bytes_to_go(py_bytes: bytes):
    c_bytes = ffi.new("char[]", py_bytes)
    go_slice = ffi.new("GoSlice *", [c_bytes, len(py_bytes), len(py_bytes)])
    return (go_slice[0], c_bytes)


def py_str_to_go(py_str: str):
    str_bytes = py_str.encode()
    c_str = ffi.new("char[]", str_bytes)
    go_str = ffi.new("_GoString_ *", [c_str, len(str_bytes)])
    return (go_str[0], c_str)


def go_str_to_py(go_str: str):
    str_len = go_str.n
    if str_len > 0:
        return ffi.unpack(go_str.p, go_str.n).decode()
    return ""


def init_go_slice():
    return ffi.new("GoSlice *")


def go_str_slice_to_py(str_slice) -> List[str]:
    slice_len = str_slice.len
    char_arr = ffi.cast("char **", str_slice.data)
    return [ffi.string(char_arr[i]).decode() for i in range(slice_len)]


def go_bool_to_py(go_bool: bool):
    return go_bool == 1


def go_guess_to_py(guess) -> Guess:
    return Guess(go_str_to_py(guess.r0), go_bool_to_py(guess.r1))


py_to_go = {
    str: py_str_to_go,
    bytes: py_bytes_to_go,
}


go_to_py = {
    str: go_str_to_py,
    bool: go_bool_to_py,
    Guess: go_guess_to_py,
}


def transform_types(in_types: Sequence[Hashable], out_type: Hashable):
    def decorator(fn):
        @wraps(fn)
        def inner(*args):
            args_transformed = [py_to_go[type_](arg) for type_, arg in zip(in_types, args)]
            return go_to_py[out_type](fn(*(arg[0] for arg in args_transformed)))
        return inner
    return decorator


def transform_types_ret_str_slice(in_types: Sequence[Hashable]):
    def decorator(fn):
        @wraps(fn)
        def inner(*args):
            ret_slice = init_go_slice()
            args_transformed = [py_to_go[type_](arg) for type_, arg in zip(in_types, args)]
            fn(*(arg[0] for arg in args_transformed), ret_slice)
            return go_str_slice_to_py(ret_slice)
        return inner
    return decorator