Source code for labgrid.util.agent

#!/usr/bin/env python3

import json
import os
import signal
import sys
import base64
import types

[docs]def b2s(b): return base64.b85encode(b).decode('ascii')
[docs]def s2b(s): return base64.b85decode(s.encode('ascii'))
[docs]class Agent:
[docs] def __init__(self): self.methods = {} self.register('load', self.load) self.register('list', self.list) # use real stdin/stdout self.stdin = sys.stdin self.stdout = sys.stdout # use stderr for normal prints sys.stdout = sys.stderr
[docs] def send(self, data): self.stdout.write(json.dumps(data)+'\n') self.stdout.flush()
[docs] def register(self, name, func): assert name not in self.methods self.methods[name] = func
[docs] def load(self, name, source): module = types.ModuleType(name) exec(compile(source, '<loaded {}>'.format(name), 'exec'), module.__dict__) for k, v in module.methods.items(): self.register('{}.{}'.format(name, k), v)
[docs] def list(self): return list(self.methods.keys())
[docs] def run(self): for line in self.stdin: if not line: continue try: request = json.loads(line) except json.JSONDecodeError: self.send({'error': 'request parsing failed for {}'.format(repr(line))}) break if request.get('close', False): break name = request['method'] args = request['args'] kwargs = request['kwargs'] try: response = self.methods[name](*args, **kwargs) self.send({'result': response}) except Exception as e: # pylint: disable=broad-except import traceback try: tb = [list(x) for x in traceback.extract_tb(sys.exc_info()[2])] except: tb = None self.send({'exception': repr(e), 'tb': tb})
[docs]def handle_test(*args, **kwargs): # pylint: disable=unused-argument return args[::-1]
[docs]def handle_error(message): raise ValueError(message)
[docs]def handle_usbtmc(index, cmd, read=False): assert isinstance(index, int) cmd = s2b(cmd) fd = os.open('/dev/usbtmc{}'.format(index), os.O_RDWR) os.write(fd, cmd) if not read: os.close(fd) return None data = [] while True: data.append(os.read(fd, 4096)) if len(data[-1]) < 4096: break os.close(fd) return b2s(b''.join(data))
[docs]def main(): signal.signal(signal.SIGINT, signal.SIG_IGN) a = Agent() a.register('test', handle_test) a.register('error', handle_error) a.register('usbtmc', handle_usbtmc) a.run()
if __name__ == "__main__": main()