rtt_command.py

#!/usr/bin/env python3

import argparse
import sys
import time
import logging
import os

# Add pynrfjprog to Python path
pynrfjprog_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))), 
                              'tools', 'nordic-cli', 'nrf-command-line-tools', 'python')
sys.path.insert(0, pynrfjprog_path)

from pynrfjprog import LowLevel
from pynrfjprog.APIError import APIError
from west.commands import WestCommand

class RttCommand(WestCommand):
    def __init__(self):
        super().__init__(
            'rtt',
            'Real-Time Transfer (RTT) debugging tool',
            'Connect to a Nordic device and read RTT data.')
        self.logger = logging.getLogger('west.rtt')
        self.logger.setLevel(logging.INFO)
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(message)s')
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)

    def do_add_parser(self, parser_adder):
        parser = parser_adder.add_parser(
            self.name,
            help=self.help,
            description=self.description)
        parser.add_argument('-s', '--serial', type=int, help='Serial number of the device to connect to')
        parser.add_argument('-c', '--channel', type=int, default=0, help='RTT channel to use (default: 0)')
        parser.add_argument('-b', '--baudrate', type=int, default=115200, help='Baudrate for RTT (default: 115200)')
        parser.add_argument('-t', '--timeout', type=float, default=5.0, help='Timeout in seconds for RTT operations (default: 5.0)')
        return parser

    def do_run(self, args, unknown_args):
        try:
            with LowLevel.API() as api:
                if args.serial:
                    api.connect_to_emu_with_snr(args.serial)
                else:
                    api.connect_to_emu_without_snr()

                family = api.read_device_family()
                self.logger.info(f"Connected to device: {family}")

                api.rtt_start()
                self.logger.info("RTT started")

                start_time = time.time()
                while not api.rtt_is_control_block_found():
                    if time.time() - start_time > args.timeout:
                        self.logger.error("Timeout waiting for RTT control block")
                        return 1
                    time.sleep(0.1)

                down_channel_count, up_channel_count = api.rtt_read_channel_count()
                self.logger.info(f"Found {up_channel_count} up channels and {down_channel_count} down channels")

                for idx in range(up_channel_count):
                    name, size = api.rtt_read_channel_info(idx, 'UP_DIRECTION')
                    self.logger.info(f"Up channel {idx}: {name} (buffer size: {size})")

                for idx in range(down_channel_count):
                    name, size = api.rtt_read_channel_info(idx, 'DOWN_DIRECTION')
                    self.logger.info(f"Down channel {idx}: {name} (buffer size: {size})")

                channel_index = args.channel
                self.logger.info(f"Using channel {channel_index}")

                self.logger.info("Starting RTT read (press Ctrl+C to stop)...")
                try:
                    while True:
                        data = api.rtt_read(channel_index=channel_index, length=4096)
                        if data:
                            sys.stdout.write(data)
                            sys.stdout.flush()
                        time.sleep(0.0001)
                except KeyboardInterrupt:
                    self.logger.info("\nStopping RTT...")
                    api.rtt_stop()
                    return 0

        except APIError as e:
            self.logger.error(f"Error: {str(e)}")
            return 1
        except Exception as e:
            self.logger.error(f"Unexpected error: {str(e)}")
            return 1
        finally:
            try:
                api.rtt_stop()
                api.disconnect_from_emu()
            except:
                pass

        return 0