refactor: rework pb2py generator

This also includes the capability to build Rust protobuf blobs.

core/site_scons/site_tools/micropython/__init__.py

@@ -25,9 +25,13 @@ def generate(env):
         source_name = source.replace(env['source_dir'], '')
         # replace "utils.BITCOIN_ONLY" with literal constant (True/False)
         # so the compiler can optimize out the things we don't want
-        btc_only = 'True' if env['bitcoin_only'] == '1' else 'False'
+        btc_only = env['bitcoin_only'] == '1'
         interim = "%s.i" % target[:-4]  # replace .mpy with .i
-        return '$SED "s:utils\\.BITCOIN_ONLY:%s:g" %s > %s && $MPY_CROSS -o %s -s %s %s' % (btc_only, source, interim, target, source_name, interim)
+        sed_scripts = " ".join([
+            f"-e 's/utils\.BITCOIN_ONLY/{btc_only}/g'",
+            "-e 's/if TYPE_CHECKING/if False/'",
+        ])
+        return f'$SED {sed_scripts} {source} > {interim} && $MPY_CROSS -o {target} -s {source_name} {interim}'
 
     env['BUILDERS']['FrozenModule'] = SCons.Builder.Builder(
         generator=generate_frozen_module,

core/tools/build_templates

@@ -4,13 +4,11 @@ set -e
 CWD=`dirname "$0"`
 RENDER="$CWD/../vendor/trezor-common/tools/cointool.py render"
 
-FIND_TEMPLATES="find $CWD/../src -name *.mako"
+FIND_TEMPLATES="find $CWD/../src -name *.mako -not -name _proto*"
 
 check_results() {
     CHECK_FAIL=0
     for filename in $($FIND_TEMPLATES); do
-        # ignore resources.py
-        if echo $filename | grep -q "resources.py.mako$"; then continue; fi
         TMP=`mktemp`
         TARGET="${filename%%.mako}"
         $RENDER "$filename" -o $TMP

tools/build_protobuf

@@ -51,55 +51,41 @@ PYTHON_MESSAGES_IGNORE=""
 RETURN=0
 
 do_rebuild() {
-    # rebuild protobuf in specified directory
+    local FILE_OR_DIR="$1"
-    local DESTDIR="$1"
     shift
-    local SOURCES="$1"
+    local OUTPUT="$1"
-    shift
-    local IGNORE="$1"
     shift
-    local APPLY_BITCOIN_ONLY="$1"
+    local SOURCES="$1"
     shift
 
-    mkdir -p "$DESTDIR"
+    if [ "$FILE_OR_DIR" == file ]; then
-    rm -f "$DESTDIR"/[A-Z]*.py
+        local param="--outfile"
-
+    else
-    # note $SOURCES is unquoted - we want wildcard expansion and multiple args
+        local param="--python-outdir"
-    $PROTOB/pb2py "$@" -o "$DESTDIR" $SOURCES
-
-    # TODO: make this less hackish
-    # maybe introduce attribute "altcoin" in protobuf?
-    if [ "$APPLY_BITCOIN_ONLY" == "TRUE" ]; then
-        sed -i "3ifrom trezor import utils\n" "$DESTDIR"/Capability.py
-        sed -i "3ifrom trezor import utils\n" "$DESTDIR"/MessageType.py
-        sed -i "/^EthereumGetPublicKey/iif not utils.BITCOIN_ONLY:" "$DESTDIR"/MessageType.py
-        for altcoin in Ethereum NEM Lisk Tezos Stellar Cardano Ripple Monero DebugMonero Eos Binance WebAuthn; do
-            sed -i "s:^$altcoin:    $altcoin:" "$DESTDIR"/Capability.py
-            sed -i "s:^$altcoin:    $altcoin:" "$DESTDIR"/MessageType.py
-        done
-        sed -i "/^Bitcoin_like/iif not utils.BITCOIN_ONLY:" "$DESTDIR"/Capability.py
-        sed -i "/^EOS/iif not utils.BITCOIN_ONLY:" "$DESTDIR"/Capability.py
-        for feature in Bitcoin_like EOS U2F; do
-            sed -i "s:^$feature:    $feature:" "$DESTDIR"/Capability.py
-        done
     fi
-    # ENDTODO
 
-    # delete unused messages
+    # note $SOURCES is unquoted - we want wildcard expansion and multiple args
-    for F in $IGNORE; do
+    $PROTOB/pb2py "$@" $param="$OUTPUT" $SOURCES
-        rm -f "$DESTDIR"/"$F".py
-    done
 }
 
 do_check() {
     # rebuild protobuf in tmpdir and check result against specified directory
     local TMPDIR=$(mktemp -d proto-check.XXXXXX)
-    local DESTDIR="$1"
+
+    local FILE_OR_DIR="$1"
+    shift
+    local OUTPUT="$1"
     shift
 
-    cp -rT "$DESTDIR" "$TMPDIR"
+    if [ "$FILE_OR_DIR" == file ]; then
-    do_rebuild "$TMPDIR" "$@"
+        local TMPDEST="$TMPDIR/testfile"
-    DIFF=$(diff -ur --exclude __pycache__ "$DESTDIR" "$TMPDIR")
+    else
+        cp -rT "$OUTPUT" "$TMPDIR"
+        local TMPDEST="$TMPDIR"
+    fi
+
+    do_rebuild "$FILE_OR_DIR" "$TMPDEST" "$@"
+    DIFF=$(diff -ur --exclude __pycache__ "$OUTPUT" "$TMPDEST")
     rm -r "$TMPDIR"
     if [ -n "$DIFF" ]; then
         echo "$DIFF"
@@ -113,7 +99,12 @@ else
     func=do_rebuild
 fi
 
-$func core/src/trezor/messages "$CORE_PROTOBUF_SOURCES" "$CORE_MESSAGES_IGNORE" TRUE --no-init-py
+$func dir core/src/trezor/enums "$CORE_PROTOBUF_SOURCES"
-$func python/src/trezorlib/messages "$PYTHON_PROTOBUF_SOURCES" "$PYTHON_MESSAGES_IGNORE" FALSE --include-deprecated -P ..protobuf
+$func file core/src/trezor/enums/__init__.py "$CORE_PROTOBUF_SOURCES" --template=core/src/trezor/enums/_proto_init.mako
+$func file core/src/trezor/messages.py "$CORE_PROTOBUF_SOURCES" --template=core/src/trezor/_proto_messages.mako
+
+$func file python/src/trezorlib/messages.py "$PYTHON_PROTOBUF_SOURCES" \
+    --template=python/src/trezorlib/_proto_messages.mako \
+    --include-deprecated
 
 exit $RETURN