Add option to terminate TCP connections after some time

This commit is contained in:
jaseg 2019-04-04 19:55:25 +09:00
parent 8c5689171c
commit e3e55e031e

487
clippy.py
View file

@ -24,268 +24,275 @@ HOST, PORT = "172.23.42.29",2342
CMD_LED_DRAW = 18 CMD_LED_DRAW = 18
def resize_image(img, size, blackbg=True): def resize_image(img, size, blackbg=True):
tw, th = size tw, th = size
w, h = img.size w, h = img.size
a, b = w/tw, h/th a, b = w/tw, h/th
f = 1/max(a, b) f = 1/max(a, b)
pos = int((tw-w*f)/2), int((th-h*f)/2) pos = int((tw-w*f)/2), int((th-h*f)/2)
buf = Image.new('RGBA', (tw, th)) buf = Image.new('RGBA', (tw, th))
buf.paste(img.resize((int(w*f), int(h*f))).convert('RGBA'), pos) buf.paste(img.resize((int(w*f), int(h*f))).convert('RGBA'), pos)
if blackbg: if blackbg:
buf2 = Image.new('RGBA', (tw, th), (0, 0, 0, 255)) buf2 = Image.new('RGBA', (tw, th), (0, 0, 0, 255))
return Image.alpha_composite(buf2, buf) return Image.alpha_composite(buf2, buf)
else: else:
return buf return buf
class Display: class Display:
def __init__(self): def __init__(self):
self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
self.size = 56*8, 8*20 self.size = 56*8, 8*20
def sendframe(self, frame): def sendframe(self, frame):
pl = struct.pack('!HHHHH', CMD_LED_DRAW, 0, 0, 0x627a, 0) + bz2.compress(frame) pl = struct.pack('!HHHHH', CMD_LED_DRAW, 0, 0, 0x627a, 0) + bz2.compress(frame)
self.sock.sendto(pl, (HOST, PORT)) self.sock.sendto(pl, (HOST, PORT))
# for i in range(100): # for i in range(100):
# time.sleep(0.0001) # time.sleep(0.0001)
self.sock.sendto(pl, (HOST, PORT)) self.sock.sendto(pl, (HOST, PORT))
@staticmethod @staticmethod
def do_gamma(im, gamma): def do_gamma(im, gamma):
"""Fast gamma correction with PIL's image.point() method""" """Fast gamma correction with PIL's image.point() method"""
invert_gamma = 1.0/gamma invert_gamma = 1.0/gamma
lut = [pow(x/255., invert_gamma) * 255 for x in range(256)] lut = [pow(x/255., invert_gamma) * 255 for x in range(256)]
lut = lut*4 # need one set of data for each band for RGBA lut = lut*4 # need one set of data for each band for RGBA
im = im.point(lut) im = im.point(lut)
return im return im
@staticmethod @staticmethod
def encode_image(img, displaysize): def encode_image(img, displaysize):
return np.frombuffer(Display.do_gamma(resize_image(img, displaysize), 0.5).convert('1').tobytes(), dtype='1b') return np.frombuffer(Display.do_gamma(resize_image(img, displaysize), 0.5).convert('1').tobytes(), dtype='1b')
def weightedChoice(choices, default=None): def weightedChoice(choices, default=None):
acc = 0 acc = 0
r = random.random() r = random.random()
for weight, choice in choices: for weight, choice in choices:
if r < (acc + weight): if r < (acc + weight):
return choice return choice
acc += weight acc += weight
return default return default
class Agent: class Agent:
def __init__(self, path: 'pathlib.Path'): def __init__(self, path: 'pathlib.Path'):
self.config = json.loads((path / 'agent.json').read_text()) self.config = json.loads((path / 'agent.json').read_text())
for ani in self.config['animations'].values(): for ani in self.config['animations'].values():
for f in ani['frames']: for f in ani['frames']:
branching, exitBranch = f.get('branching'), f.get('exitBranch') branching, exitBranch = f.get('branching'), f.get('exitBranch')
if 'exitBranch' in f: if 'exitBranch' in f:
f['next'] = lambda f, idx: f['exitBranch'] f['next'] = lambda f, idx: f['exitBranch']
elif 'branching' in f: elif 'branching' in f:
f['next'] = lambda f, idx: weightedChoice( f['next'] = lambda f, idx: weightedChoice(
[ (b['weight']/100, b['frameIndex']) for b in f['branching']['branches'] ] [ (b['weight']/100, b['frameIndex']) for b in f['branching']['branches'] ]
, default=idx+1) , default=idx+1)
else: else:
f['next'] = lambda f, idx: idx+1 f['next'] = lambda f, idx: idx+1
self.picmap = Image.open(path / 'map.png') self.picmap = Image.open(path / 'map.png')
self.path = path self.path = path
def __call__(self, action, sleep=True): def __call__(self, action, sleep=True):
for frame in self._animate(action): for frame in self._animate(action):
# print('frame:', frame) # print('frame:', frame)
if 'images_encoded' in frame: # some frames contain branch info and sound, but no images if 'images_encoded' in frame: # some frames contain branch info and sound, but no images
yield frame['images_encoded'] yield frame['images_encoded']
if sleep: if sleep:
time.sleep(frame['duration']/1000) time.sleep(frame['duration']/1000)
def precalculate_images(self, pf, dsp, termsize): def precalculate_images(self, pf, dsp, termsize):
print('\033[93mPrecalculating images\033[0m', flush=True) print('\033[93mPrecalculating images\033[0m', flush=True)
total = sum(1 for ani in self.config['animations'].values() for f in ani['frames'] if 'images' in f) total = sum(1 for ani in self.config['animations'].values() for f in ani['frames'] if 'images' in f)
i = 0 i = 0
for ani in self.config['animations'].values(): for ani in self.config['animations'].values():
for f in ani['frames']: for f in ani['frames']:
if 'images' in f: if 'images' in f:
print(('(\033[38;5;245m{: '+str(1+int(math.log10(total)))+'}/{}\033[0m)').format(i, total), end='', flush=True) print(('(\033[38;5;245m{: '+str(1+int(math.log10(total)))+'}/{}\033[0m)').format(i, total), end='', flush=True)
i += 1 i += 1
f['images_encoded'] = self._precalculate_one_image(tuple(f['images'][0]), pf, dsp, termsize) f['images_encoded'] = self._precalculate_one_image(tuple(f['images'][0]), pf, dsp, termsize)
print(flush=True) print(flush=True)
print('\033[93mdone.\033[0m', flush=True) print('\033[93mdone.\033[0m', flush=True)
self._precalculate_one_image.cache_clear() self._precalculate_one_image.cache_clear()
@functools.lru_cache(maxsize=None) @functools.lru_cache(maxsize=None)
def _precalculate_one_image(self, coords, pf, dsp, termsize): def _precalculate_one_image(self, coords, pf, dsp, termsize):
img = self._get_image(*coords) img = self._get_image(*coords)
return ( pf.encode_image(img) if pf else None, return ( pf.encode_image(img) if pf else None,
dsp.encode_image(img, dsp.size) if dsp else None, dsp.encode_image(img, dsp.size) if dsp else None,
pixelterm.termify_pixels(resize_image(img, termsize)) if termsize else None ) pixelterm.termify_pixels(resize_image(img, termsize)) if termsize else None )
def _animate(self, action): def _animate(self, action):
anim, idx = self.config['animations'][action]['frames'], 0 anim, idx = self.config['animations'][action]['frames'], 0
while idx < len(anim): while idx < len(anim):
yield anim[idx] yield anim[idx]
idx = anim[idx]['next'](anim[idx], idx) idx = anim[idx]['next'](anim[idx], idx)
def _get_image(self, x, y): def _get_image(self, x, y):
print('\033[38;5;96mcropbox:\033[0m {:04} {:04} {:04} {:04} \033[38;5;96mmap:\033[0m {:04} {:04}'.format( print('\033[38;5;96mcropbox:\033[0m {:04} {:04} {:04} {:04} \033[38;5;96mmap:\033[0m {:04} {:04}'.format(
x, y, *self.config['framesize'], *self.picmap.size), end='', flush=True) x, y, *self.config['framesize'], *self.picmap.size), end='', flush=True)
tw, th = self.config['framesize'] tw, th = self.config['framesize']
return self.picmap.crop((x, y, x+tw, y+th)) return self.picmap.crop((x, y, x+tw, y+th))
@property @property
def animations(self): def animations(self):
return list(self.config['animations'].keys()) return list(self.config['animations'].keys())
if __name__ == '__main__': if __name__ == '__main__':
import argparse, pathlib, sys import argparse, pathlib, sys
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('-l', '--list', action='store_true') parser.add_argument('-l', '--list', action='store_true')
parser.add_argument('-a', '--agent', default='Clippy') parser.add_argument('-a', '--agent', default='Clippy')
parser.add_argument('-e', '--endless', action='store_true') parser.add_argument('-e', '--endless', action='store_true')
parser.add_argument('-d', '--display', action='store_true') parser.add_argument('-d', '--display', action='store_true')
parser.add_argument('-i', '--interactive', action='store_true') parser.add_argument('-i', '--interactive', action='store_true')
parser.add_argument('-w', '--wait', type=int, default=120) parser.add_argument('-w', '--wait', type=int, default=120)
parser.add_argument('-p', '--pixelflut', type=str) parser.add_argument('-p', '--pixelflut', type=str)
parser.add_argument('-t', '--terminal', action='store_true') parser.add_argument('-t', '--terminal', action='store_true')
parser.add_argument('-x', '--termsize', type=str) parser.add_argument('-x', '--termsize', type=str)
parser.add_argument('-s', '--socket', action='store_true') parser.add_argument('-s', '--socket', action='store_true', help='Listen on TCP socket (telnet-compatible)')
parser.add_argument('-n', '--nosleep', action='store_true') parser.add_argument('-k', '--kill-after', type=int, default=60, help='Kill TCP connections after {kill_after} seconds')
parser.add_argument('-b', '--bind', type=str, default='0.0.0.0:2342') parser.add_argument('-n', '--nosleep', action='store_true')
parser.add_argument('action', default='Greeting', nargs='?') parser.add_argument('-b', '--bind', type=str, default='0.0.0.0:2342')
args = parser.parse_args() parser.add_argument('action', default='Greeting', nargs='?')
args = parser.parse_args()
agent_paths = [] agent_paths = []
for agent in args.agent.split(','): for agent in args.agent.split(','):
agent_path = pathlib.Path('agents') / agent agent_path = pathlib.Path('agents') / agent
if not agent_path.is_dir(): if not agent_path.is_dir():
print('Agent "{}" not found. Exiting.'.format(agent), flush=True) print('Agent "{}" not found. Exiting.'.format(agent), flush=True)
sys.exit(1) sys.exit(1)
agent_paths.append(agent_path) agent_paths.append(agent_path)
if args.list: if args.list:
print('\n'.join(Agent(agent_path).animations), flush=True) print('\n'.join(Agent(agent_path).animations), flush=True)
sys.exit(0) sys.exit(0)
dsp = Display() if args.display else None dsp = Display() if args.display else None
if args.pixelflut: if args.pixelflut:
target, *params = args.pixelflut.split('@') target, *params = args.pixelflut.split('@')
host, port = target.split(':') host, port = target.split(':')
port = int(port) port = int(port)
x, y, *_r = params[0].split(',') if params else (0, 0, None) x, y, *_r = params[0].split(',') if params else (0, 0, None)
w, h, reps = _r if _r else (320, 240) w, h, reps = _r if _r else (320, 240)
x, y, w, h, reps = map(int, (x, y, w, h, reps)) x, y, w, h, reps = map(int, (x, y, w, h, reps))
pf = pxf.Pixelflut(host, port, x, y, w, h, reps) if args.pixelflut else None pf = pxf.Pixelflut(host, port, x, y, w, h, reps) if args.pixelflut else None
else: else:
pf = None pf = None
agents = [] agents = []
for path in agent_paths: for path in agent_paths:
agent = Agent(path) agent = Agent(path)
if args.socket: if args.socket:
tx, ty = (args.termsize or '60x30').split('x') tx, ty = (args.termsize or '60x30').split('x')
tx, ty = int(tx), int(ty) tx, ty = int(tx), int(ty)
elif args.terminal: elif args.terminal:
tx, ty = args.termsize.split('x') or os.get_terminal_size() tx, ty = args.termsize.split('x') or os.get_terminal_size()
tx, ty = int(tx), int(ty) tx, ty = int(tx), int(ty)
termsize = (tx, ty*2) if args.terminal or args.socket else None termsize = (tx, ty*2) if args.terminal or args.socket else None
agent.precalculate_images(pf, dsp, termsize) agent.precalculate_images(pf, dsp, termsize)
agents.append(agent) agents.append(agent)
runlock = threading.Lock() runlock = threading.Lock()
ts = time.time() ts = time.time()
if args.interactive: if args.interactive:
from tkinter import * from tkinter import *
def recalc_size(delta): def recalc_size(delta):
global runlock global runlock
with runlock: with runlock:
print('resetting', flush=True) print('resetting', flush=True)
pf.reset_images() pf.reset_images()
pf.w += delta pf.w += delta
pf.h += delta pf.h += delta
print('recalcing', flush=True) print('recalcing', flush=True)
for agent in Agents: for agent in Agents:
agent.precalculate_images(pf, dsp, termsize) agent.precalculate_images(pf, dsp, termsize)
def keyfunc(ev): def keyfunc(ev):
global ts global ts
ch = ev.char ch = ev.char
if ch == '+': if ch == '+':
recalc_size(50) recalc_size(50)
elif ch == '-': elif ch == '-':
recalc_size(-50) recalc_size(-50)
if ch == 'w': if ch == 'w':
pf.y -= 10 pf.y -= 10
elif ch == 'a': elif ch == 'a':
pf.x -= 10 pf.x -= 10
elif ch == 's': elif ch == 's':
pf.y += 10 pf.y += 10
elif ch == 'd': elif ch == 'd':
pf.x += 10 pf.x += 10
elif ch == 'e': elif ch == 'e':
pf.reps += 1 pf.reps += 1
elif ch == 'q': elif ch == 'q':
if pf.reps > 1: if pf.reps > 1:
pf.reps -= 1 pf.reps -= 1
elif ch == 'n': elif ch == 'n':
ts = time.time() - args.wait - 1 ts = time.time() - args.wait - 1
def tkrun(): def tkrun():
tkr = Tk() tkr = Tk()
tkf = Frame(tkr, width=100, height=100) tkf = Frame(tkr, width=100, height=100)
tkf.bind('<Key>', keyfunc) tkf.bind('<Key>', keyfunc)
tkf.pack() tkf.pack()
tkf.focus_set() tkf.focus_set()
tkr.mainloop() tkr.mainloop()
tkrunner = threading.Thread(target=tkrun, daemon=True) tkrunner = threading.Thread(target=tkrun, daemon=True)
tkrunner.start() tkrunner.start()
if args.socket: if args.socket:
import socketserver import socketserver
class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer): class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
pass pass
class ClippyRequestHandler(socketserver.BaseRequestHandler): class ClippyRequestHandler(socketserver.BaseRequestHandler):
def handle(self): def handle(self):
with contextlib.suppress(BrokenPipeError): with contextlib.suppress(BrokenPipeError):
agent = random.choice(agents) start = time.time()
while True: srcaddr, srcport = self.client_address
action = random.choice(agent.animations) print(f'Connection from {srcaddr}:{srcport}')
print('[\033[38;5;245m{}\033[0m] Playing: {}'.format(self.client_address[0], action), flush=True)
for _img_pf, _img_dsp, img_term in agent(action): agent = random.choice(agents)
self.request.sendall(b'\033[H'+img_term.encode()) while True:
host, port = args.bind.split(':') action = random.choice(agent.animations)
port = int(port) #print('[\033[38;5;245m{}\033[0m] Playing: {}'.format(self.client_address[0], action), flush=True)
server = ThreadedTCPServer((host, port), ClippyRequestHandler) for _img_pf, _img_dsp, img_term in agent(action):
server.serve_forever() if time.time() - start > args.kill_after:
elif args.endless: return
while True: self.request.sendall(b'\033[H'+img_term.encode())
print('Starting', ts, flush=True) host, port = args.bind.split(':')
for agent in agents: port = int(port)
while time.time() - ts < args.wait: server = ThreadedTCPServer((host, port), ClippyRequestHandler)
if random.random() > 0.2: server.serve_forever()
action = random.choice(agent.animations) elif args.endless:
print('Playing:', action, flush=True) while True:
for img_pf, img_dsp, img_term in agent(action, not args.nosleep): print('Starting', ts, flush=True)
with runlock: for agent in agents:
if args.terminal: while time.time() - ts < args.wait:
print('\033[H'+img_term, flush=True) if random.random() > 0.2:
if args.display: action = random.choice(agent.animations)
dsp.sendframe(img_dsp) print('Playing:', action, flush=True)
if args.pixelflut: for img_pf, img_dsp, img_term in agent(action, not args.nosleep):
pf.sendframe(img_pf) with runlock:
if time.time() - ts > args.wait: if args.terminal:
print('Force-advance', ts, flush=True) print('\033[H'+img_term, flush=True)
break if args.display:
if not args.nosleep: dsp.sendframe(img_dsp)
time.sleep(1) if args.pixelflut:
print('Advancing', ts, flush=True) pf.sendframe(img_pf)
ts = time.time() if time.time() - ts > args.wait:
else: print('Force-advance', ts, flush=True)
for img_pf, img_dsp, img_term in agents[0](args.action, not args.nosleep): break
if args.terminal: if not args.nosleep:
print(img_term, flush=True) #pixelterm.termify_pixels( time.sleep(1)
#resize_image(img, termsize))) print('Advancing', ts, flush=True)
if args.display: ts = time.time()
dsp.sendframe(img_dsp) else:
if args.pixelflut: for img_pf, img_dsp, img_term in agents[0](args.action, not args.nosleep):
pf.sendframe(img_pf) if args.terminal:
print(img_term, flush=True) #pixelterm.termify_pixels(
#resize_image(img, termsize)))
if args.display:
dsp.sendframe(img_dsp)
if args.pixelflut:
pf.sendframe(img_pf)