You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

1182 lines
38 KiB

4 years ago
  1. # this module requires Python 2.6+
  2. from __future__ import with_statement
  3. from contextlib import contextmanager
  4. from operator import attrgetter
  5. from gettext import gettext as _
  6. import imp
  7. import inspect
  8. import os
  9. import sys
  10. import cmd
  11. import shlex
  12. import subprocess
  13. import argparse
  14. import itertools
  15. import traceback
  16. import multiprocessing
  17. import signal
  18. import threading
  19. import plac_core
  20. if sys.version < '3':
  21. def exec_(_code_, _globs_=None, _locs_=None):
  22. if _globs_ is None:
  23. frame = sys._getframe(1)
  24. _globs_ = frame.f_globals
  25. if _locs_ is None:
  26. _locs_ = frame.f_locals
  27. del frame
  28. elif _locs_ is None:
  29. _locs_ = _globs_
  30. exec("""exec _code_ in _globs_, _locs_""")
  31. exec('''
  32. def raise_(tp, value=None, tb=None):
  33. raise tp, value, tb
  34. ''')
  35. else:
  36. exec_ = eval('exec')
  37. def raise_(tp, value=None, tb=None):
  38. """
  39. A function that matches the Python 2.x ``raise`` statement. This
  40. allows re-raising exceptions with the cls value and traceback on
  41. Python 2 and 3.
  42. """
  43. if value is not None and isinstance(tp, Exception):
  44. raise TypeError("instance exception may not have a separate value")
  45. if value is not None:
  46. exc = tp(value)
  47. else:
  48. exc = tp
  49. if exc.__traceback__ is not tb:
  50. raise exc.with_traceback(tb)
  51. raise exc
  52. def decode(val):
  53. """
  54. Decode an object assuming the encoding is UTF-8.
  55. """
  56. try:
  57. # assume it is an encoded bytes object
  58. return val.decode('utf-8')
  59. except AttributeError:
  60. # it was an already decoded unicode object
  61. return val
  62. # ############################ generic utils ############################### #
  63. @contextmanager
  64. def stdout(fileobj):
  65. "usage: with stdout(file('out.txt', 'a')): do_something()"
  66. orig_stdout = sys.stdout
  67. sys.stdout = fileobj
  68. try:
  69. yield
  70. finally:
  71. sys.stdout = orig_stdout
  72. def write(x):
  73. "Write str(x) on stdout and flush, no newline added"
  74. sys.stdout.write(str(x))
  75. sys.stdout.flush()
  76. def gen_val(value):
  77. "Return a generator object with a single element"
  78. yield value
  79. def gen_exc(etype, exc, tb):
  80. "Return a generator object raising an exception"
  81. raise_(etype, exc, tb)
  82. yield
  83. def less(text):
  84. "Send a text to less via a pipe"
  85. # -c clear the screen before starting less
  86. po = subprocess.Popen(['less', '-c'], stdin=subprocess.PIPE)
  87. try:
  88. po.stdin.write(text)
  89. except IOError:
  90. pass
  91. po.stdin.close()
  92. po.wait()
  93. use_less = (sys.platform != 'win32') # unices
  94. class TerminatedProcess(Exception):
  95. pass
  96. def terminatedProcess(signum, frame):
  97. raise TerminatedProcess
  98. # ########################## readline support ############################ #
  99. def read_line(stdin, prompt=''):
  100. "Read a line from stdin, using readline when possible"
  101. if isinstance(stdin, ReadlineInput):
  102. return stdin.readline(prompt)
  103. else:
  104. write(prompt)
  105. return stdin.readline()
  106. def read_long_line(stdin, terminator):
  107. """
  108. Read multiple lines from stdin until the terminator character is found,
  109. then yield a single space-separated long line.
  110. """
  111. while True:
  112. lines = []
  113. while True:
  114. line = stdin.readline() # ends with \n
  115. if not line: # EOF
  116. return
  117. line = line.strip()
  118. if not line:
  119. continue
  120. elif line[-1] == terminator:
  121. lines.append(line[:-1])
  122. break
  123. else:
  124. lines.append(line)
  125. yield ' '.join(lines)
  126. class ReadlineInput(object):
  127. """
  128. An iterable with a .readline method reading from stdin.
  129. """
  130. def __init__(self, completions, case_sensitive=True, histfile=None):
  131. self.completions = completions
  132. self.case_sensitive = case_sensitive
  133. self.histfile = histfile
  134. if not case_sensitive:
  135. self.completions = [c.upper() for c in completions]
  136. import readline
  137. self.rl = readline
  138. readline.parse_and_bind("tab: complete")
  139. readline.set_completer(self.complete)
  140. def __enter__(self):
  141. self.old_completer = self.rl.get_completer()
  142. try:
  143. if self.histfile:
  144. self.rl.read_history_file(self.histfile)
  145. except IOError: # the first time
  146. pass
  147. return self
  148. def __exit__(self, etype, exc, tb):
  149. self.rl.set_completer(self.old_completer)
  150. if self.histfile:
  151. self.rl.write_history_file(self.histfile)
  152. def complete(self, kw, state):
  153. # state is 0, 1, 2, ... and increases by hitting TAB
  154. if not self.case_sensitive:
  155. kw = kw.upper()
  156. try:
  157. return [k for k in self.completions if k.startswith(kw)][state]
  158. except IndexError: # no completions
  159. return # exit
  160. def readline(self, prompt=''):
  161. try:
  162. return raw_input(prompt) + '\n'
  163. except EOFError:
  164. return ''
  165. def __iter__(self):
  166. return iter(self.readline, '')
  167. # ################# help functionality in plac interpreters ################# #
  168. class HelpSummary(object):
  169. "Build the help summary consistently with the cmd module"
  170. @classmethod
  171. def add(cls, obj, specialcommands):
  172. p = plac_core.parser_from(obj)
  173. c = cmd.Cmd(stdout=cls())
  174. c.stdout.write('\n')
  175. c.print_topics('special commands',
  176. sorted(specialcommands), 15, 80)
  177. c.print_topics('custom commands',
  178. sorted(obj.commands), 15, 80)
  179. c.print_topics('commands run in external processes',
  180. sorted(obj.mpcommands), 15, 80)
  181. c.print_topics('threaded commands',
  182. sorted(obj.thcommands), 15, 80)
  183. p.helpsummary = str(c.stdout)
  184. def __init__(self):
  185. self._ls = []
  186. def write(self, s):
  187. self._ls.append(s)
  188. def __str__(self):
  189. return ''.join(self._ls)
  190. class PlacFormatter(argparse.RawDescriptionHelpFormatter):
  191. def _metavar_formatter(self, action, default_metavar):
  192. 'Remove special commands from the usage message'
  193. choices = action.choices or {}
  194. action.choices = dict((n, c) for n, c in choices.items()
  195. if not n.startswith('.'))
  196. return super(PlacFormatter, self)._metavar_formatter(
  197. action, default_metavar)
  198. def format_help(self):
  199. "Attached to plac_core.ArgumentParser for plac interpreters"
  200. try:
  201. return self.helpsummary
  202. except AttributeError:
  203. return super(plac_core.ArgumentParser, self).format_help()
  204. plac_core.ArgumentParser.format_help = format_help
  205. def default_help(obj, cmd=None):
  206. "The default help functionality in plac interpreters"
  207. parser = plac_core.parser_from(obj)
  208. if cmd is None:
  209. yield parser.format_help()
  210. return
  211. subp = parser.subparsers._name_parser_map.get(cmd)
  212. if subp is None:
  213. yield _('Unknown command %s' % cmd)
  214. elif getattr(obj, '_interact_', False): # in interactive mode
  215. formatter = subp._get_formatter()
  216. formatter._prog = cmd # remove the program name from the usage
  217. formatter.add_usage(
  218. subp.usage, [a for a in subp._actions if a.dest != 'help'],
  219. subp._mutually_exclusive_groups)
  220. formatter.add_text(subp.description)
  221. for action_group in subp._action_groups:
  222. formatter.start_section(action_group.title)
  223. formatter.add_text(action_group.description)
  224. formatter.add_arguments(a for a in action_group._group_actions
  225. if a.dest != 'help')
  226. formatter.end_section()
  227. yield formatter.format_help()
  228. else: # regular argparse help
  229. yield subp.format_help()
  230. # ######################## import management ############################## #
  231. try:
  232. PLACDIRS = os.environ.get('PLACPATH', '.').split(':')
  233. except:
  234. raise ValueError(_('Ill-formed PLACPATH: got %PLACPATHs') % os.environ)
  235. def partial_call(factory, arglist):
  236. "Call a container factory with the arglist and return a plac object"
  237. a = plac_core.parser_from(factory).argspec
  238. if a.defaults or a.varargs or a.varkw:
  239. raise TypeError('Interpreter.call must be invoked on '
  240. 'factories with required arguments only')
  241. required_args = ', '.join(a.args)
  242. if required_args:
  243. required_args += ',' # trailing comma
  244. code = '''def makeobj(interact, %s *args):
  245. obj = factory(%s)
  246. obj._interact_ = interact
  247. obj._args_ = args
  248. return obj\n''' % (required_args, required_args)
  249. dic = dict(factory=factory)
  250. exec_(code, dic)
  251. makeobj = dic['makeobj']
  252. makeobj.add_help = False
  253. if inspect.isclass(factory):
  254. makeobj.__annotations__ = getattr(
  255. factory.__init__, '__annotations__', {})
  256. else:
  257. makeobj.__annotations__ = getattr(
  258. factory, '__annotations__', {})
  259. makeobj.__annotations__['interact'] = (
  260. 'start interactive interpreter', 'flag', 'i')
  261. return plac_core.call(makeobj, arglist)
  262. def import_main(path, *args):
  263. """
  264. An utility to import the main function of a plac tool. It also
  265. works with command container factories.
  266. """
  267. if ':' in path: # importing a factory
  268. path, factory_name = path.split(':')
  269. else: # importing the main function
  270. factory_name = None
  271. if not os.path.isabs(path): # relative path, look at PLACDIRS
  272. for placdir in PLACDIRS:
  273. fullpath = os.path.join(placdir, path)
  274. if os.path.exists(fullpath):
  275. break
  276. else: # no break
  277. raise ImportError(_('Cannot find %s' % path))
  278. else:
  279. fullpath = path
  280. name, ext = os.path.splitext(os.path.basename(fullpath))
  281. module = imp.load_module(name, open(fullpath), fullpath, (ext, 'U', 1))
  282. if factory_name:
  283. tool = partial_call(getattr(module, factory_name), args)
  284. else:
  285. tool = module.main
  286. return tool
  287. # ############################ Task classes ############################# #
  288. # base class not instantiated directly
  289. class BaseTask(object):
  290. """
  291. A task is a wrapper over a generator object with signature
  292. Task(no, arglist, genobj), attributes
  293. .no
  294. .arglist
  295. .outlist
  296. .str
  297. .etype
  298. .exc
  299. .tb
  300. .status
  301. and methods .run and .kill.
  302. """
  303. STATES = ('SUBMITTED', 'RUNNING', 'TOBEKILLED', 'KILLED', 'FINISHED',
  304. 'ABORTED')
  305. def __init__(self, no, arglist, genobj):
  306. self.no = no
  307. self.arglist = arglist
  308. self._genobj = self._wrap(genobj)
  309. self.str, self.etype, self.exc, self.tb = '', None, None, None
  310. self.status = 'SUBMITTED'
  311. self.outlist = []
  312. def notify(self, msg):
  313. "Notifies the underlying monitor. To be implemented"
  314. def _wrap(self, genobj, stringify_tb=False):
  315. """
  316. Wrap the genobj into a generator managing the exceptions,
  317. populating the .outlist, setting the .status and yielding None.
  318. stringify_tb must be True if the traceback must be sent to a process.
  319. """
  320. self.status = 'RUNNING'
  321. try:
  322. for value in genobj:
  323. if self.status == 'TOBEKILLED': # exit from the loop
  324. raise GeneratorExit
  325. if value is not None: # add output
  326. self.outlist.append(value)
  327. self.notify(decode(value))
  328. yield
  329. except Interpreter.Exit: # wanted exit
  330. self._regular_exit()
  331. raise
  332. except (GeneratorExit, TerminatedProcess, KeyboardInterrupt):
  333. # soft termination
  334. self.status = 'KILLED'
  335. except: # unexpected exception
  336. self.etype, self.exc, tb = sys.exc_info()
  337. self.tb = ''.join(traceback.format_tb(tb)) if stringify_tb else tb
  338. self.status = 'ABORTED'
  339. else:
  340. self._regular_exit()
  341. def _regular_exit(self):
  342. self.status = 'FINISHED'
  343. try:
  344. self.str = '\n'.join(map(decode, self.outlist))
  345. except IndexError:
  346. self.str = 'no result'
  347. def run(self):
  348. "Run the inner generator"
  349. for none in self._genobj:
  350. pass
  351. def kill(self):
  352. "Set a TOBEKILLED status"
  353. self.status = 'TOBEKILLED'
  354. def wait(self):
  355. "Wait for the task to finish: to be overridden"
  356. @property
  357. def traceback(self):
  358. "Return the traceback as a (possibly empty) string"
  359. if self.tb is None:
  360. return ''
  361. elif isinstance(self.tb, (str, bytes)):
  362. return self.tb
  363. else:
  364. return ''.join(traceback.format_tb(self.tb))
  365. @property
  366. def result(self):
  367. self.wait()
  368. if self.exc:
  369. if isinstance(self.tb, (str, bytes)):
  370. raise self.etype(self.tb)
  371. else:
  372. raise_(self.etype, self.exc, self.tb or None)
  373. if not self.outlist:
  374. return None
  375. return self.outlist[-1]
  376. def __repr__(self):
  377. "String representation containing class name, number, arglist, status"
  378. return '<%s %d [%s] %s>' % (
  379. self.__class__.__name__, self.no,
  380. ' '.join(self.arglist), self.status)
  381. nulltask = BaseTask(0, [], ('skip' for dummy in (1,)))
  382. # ######################## synchronous tasks ############################## #
  383. class SynTask(BaseTask):
  384. """
  385. Synchronous task running in the interpreter loop and displaying its
  386. output as soon as available.
  387. """
  388. def __str__(self):
  389. "Return the output string or the error message"
  390. if self.etype: # there was an error
  391. return '%s: %s' % (self.etype.__name__, self.exc)
  392. else:
  393. return '\n'.join(map(str, self.outlist))
  394. class ThreadedTask(BaseTask):
  395. """
  396. A task running in a separated thread.
  397. """
  398. def __init__(self, no, arglist, genobj):
  399. BaseTask.__init__(self, no, arglist, genobj)
  400. self.thread = threading.Thread(target=super(ThreadedTask, self).run)
  401. def run(self):
  402. "Run the task into a thread"
  403. self.thread.start()
  404. def wait(self):
  405. "Block until the thread ends"
  406. self.thread.join()
  407. # ######################## multiprocessing tasks ######################### #
  408. def sharedattr(name, on_error):
  409. "Return a property to be attached to an MPTask"
  410. def get(self):
  411. try:
  412. return getattr(self.ns, name)
  413. except: # the process was killed or died hard
  414. return on_error
  415. def set(self, value):
  416. try:
  417. setattr(self.ns, name, value)
  418. except: # the process was killed or died hard
  419. pass
  420. return property(get, set)
  421. class MPTask(BaseTask):
  422. """
  423. A task running as an external process. The current implementation
  424. only works on Unix-like systems, where multiprocessing use forks.
  425. """
  426. str = sharedattr('str', '')
  427. etype = sharedattr('etype', None)
  428. exc = sharedattr('exc', None)
  429. tb = sharedattr('tb', None)
  430. status = sharedattr('status', 'ABORTED')
  431. @property
  432. def outlist(self):
  433. try:
  434. return self._outlist
  435. except: # the process died hard
  436. return []
  437. def notify(self, msg):
  438. self.man.notify_listener(self.no, msg)
  439. def __init__(self, no, arglist, genobj, manager):
  440. """
  441. The monitor has a .send method and a .man multiprocessing.Manager
  442. """
  443. self.no = no
  444. self.arglist = arglist
  445. self._genobj = self._wrap(genobj, stringify_tb=True)
  446. self.man = manager
  447. self._outlist = manager.mp.list()
  448. self.ns = manager.mp.Namespace()
  449. self.status = 'SUBMITTED'
  450. self.etype, self.exc, self.tb = None, None, None
  451. self.str = repr(self)
  452. self.proc = multiprocessing.Process(target=super(MPTask, self).run)
  453. def run(self):
  454. "Run the task into an external process"
  455. self.proc.start()
  456. def wait(self):
  457. "Block until the external process ends or is killed"
  458. self.proc.join()
  459. def kill(self):
  460. """Kill the process with a SIGTERM inducing a TerminatedProcess
  461. exception in the children"""
  462. self.proc.terminate()
  463. # ######################## Task Manager ###################### #
  464. class TaskManager(object):
  465. """
  466. Store the given commands into a task registry. Provides methods to
  467. manage the submitted tasks.
  468. """
  469. cmdprefix = '.'
  470. specialcommands = set(['.last_tb'])
  471. def __init__(self, obj):
  472. self.obj = obj
  473. self.registry = {} # {taskno : task}
  474. if obj.mpcommands or obj.thcommands:
  475. self.specialcommands.update(['.kill', '.list', '.output'])
  476. interact = getattr(obj, '_interact_', False)
  477. self.parser = plac_core.parser_from(
  478. obj, prog='' if interact else None, formatter_class=PlacFormatter)
  479. HelpSummary.add(obj, self.specialcommands)
  480. self.man = Manager() if obj.mpcommands else None
  481. signal.signal(signal.SIGTERM, terminatedProcess)
  482. def close(self):
  483. "Kill all the running tasks"
  484. for task in self.registry.values():
  485. try:
  486. if task.status == 'RUNNING':
  487. task.kill()
  488. task.wait()
  489. except: # task killed, nothing to wait
  490. pass
  491. if self.man:
  492. self.man.stop()
  493. def _get_latest(self, taskno=-1, status=None):
  494. "Get the latest submitted task from the registry"
  495. assert taskno < 0, 'You must pass a negative number'
  496. if status:
  497. tasks = [t for t in self.registry.values()
  498. if t.status == status]
  499. else:
  500. tasks = [t for t in self.registry.values()]
  501. tasks.sort(key=attrgetter('no'))
  502. if len(tasks) >= abs(taskno):
  503. return tasks[taskno]
  504. # ########################## special commands ######################## #
  505. @plac_core.annotations(
  506. taskno=('task to kill', 'positional', None, int))
  507. def kill(self, taskno=-1):
  508. 'kill the given task (-1 to kill the latest running task)'
  509. if taskno < 0:
  510. task = self._get_latest(taskno, status='RUNNING')
  511. if task is None:
  512. yield 'Nothing to kill'
  513. return
  514. elif taskno not in self.registry:
  515. yield 'Unknown task %d' % taskno
  516. return
  517. else:
  518. task = self.registry[taskno]
  519. if task.status in ('ABORTED', 'KILLED', 'FINISHED'):
  520. yield 'Already finished %s' % task
  521. return
  522. task.kill()
  523. yield task
  524. @plac_core.annotations(
  525. status=('', 'positional', None, str, BaseTask.STATES))
  526. def list(self, status='RUNNING'):
  527. 'list tasks with a given status'
  528. for task in self.registry.values():
  529. if task.status == status:
  530. yield task
  531. @plac_core.annotations(
  532. taskno=('task number', 'positional', None, int))
  533. def output(self, taskno=-1, fname=None):
  534. 'show the output of a given task (and optionally save it to a file)'
  535. if taskno < 0:
  536. task = self._get_latest(taskno)
  537. if task is None:
  538. yield 'Nothing to show'
  539. return
  540. elif taskno not in self.registry:
  541. yield 'Unknown task %d' % taskno
  542. return
  543. else:
  544. task = self.registry[taskno]
  545. outstr = '\n'.join(map(str, task.outlist))
  546. if fname:
  547. open(fname, 'w').write(outstr)
  548. yield 'saved output of %d into %s' % (taskno, fname)
  549. return
  550. yield task
  551. if len(task.outlist) > 20 and use_less:
  552. less(outstr) # has no meaning for a plac server
  553. else:
  554. yield outstr
  555. @plac_core.annotations(
  556. taskno=('task number', 'positional', None, int))
  557. def last_tb(self, taskno=-1):
  558. "show the traceback of a given task, if any"
  559. task = self._get_latest(taskno)
  560. if task:
  561. yield task.traceback
  562. else:
  563. yield 'Nothing to show'
  564. # ########################## SyncProcess ############################# #
  565. class Process(subprocess.Popen):
  566. "Start the interpreter specified by the params in a subprocess"
  567. def __init__(self, params):
  568. signal.signal(signal.SIGPIPE, signal.SIG_DFL)
  569. # to avoid broken pipe messages
  570. code = '''import plac, sys
  571. sys.argv[0] = '<%s>'
  572. plac.Interpreter(plac.import_main(*%s)).interact(prompt='i>\\n')
  573. ''' % (params[0], params)
  574. subprocess.Popen.__init__(
  575. self, [sys.executable, '-u', '-c', code],
  576. stdin=subprocess.PIPE, stdout=subprocess.PIPE)
  577. self.man = multiprocessing.Manager()
  578. def close(self):
  579. "Close stdin and stdout"
  580. self.stdin.close()
  581. self.stdout.close()
  582. self.man.shutdown()
  583. def recv(self): # char-by-char cannot work
  584. "Return the output of the subprocess, line-by-line until the prompt"
  585. lines = []
  586. while True:
  587. lines.append(self.stdout.readline())
  588. if lines[-1] == 'i>\n':
  589. out = ''.join(lines)
  590. return out[:-1] + ' ' # remove last newline
  591. def send(self, line):
  592. """Send a line (adding a newline) to the underlying subprocess
  593. and wait for the answer"""
  594. self.stdin.write(line + os.linesep)
  595. return self.recv()
  596. class StartStopObject(object):
  597. started = False
  598. def start(self):
  599. pass
  600. def stop(self):
  601. pass
  602. class Monitor(StartStopObject):
  603. """
  604. Base monitor class with methods add_listener/del_listener/notify_listener
  605. read_queue and and start/stop.
  606. """
  607. def __init__(self, name, queue=None):
  608. self.name = name
  609. self.queue = queue or multiprocessing.Queue()
  610. def add_listener(self, taskno):
  611. pass
  612. def del_listener(self, taskno):
  613. pass
  614. def notify_listener(self, taskno, msg):
  615. pass
  616. def start(self):
  617. pass
  618. def stop(self):
  619. pass
  620. def read_queue(self):
  621. pass
  622. class Manager(StartStopObject):
  623. """
  624. The plac Manager contains a multiprocessing.Manager and a set
  625. of slave monitor processes to which we can send commands. There
  626. is a manager for each interpreter with mpcommands.
  627. """
  628. def __init__(self):
  629. self.registry = {}
  630. self.started = False
  631. self.mp = None
  632. def add(self, monitor):
  633. 'Add or replace a monitor in the registry'
  634. proc = multiprocessing.Process(None, monitor.start, monitor.name)
  635. proc.queue = monitor.queue
  636. self.registry[monitor.name] = proc
  637. def delete(self, name):
  638. 'Remove a named monitor from the registry'
  639. del self.registry[name]
  640. # can be called more than once
  641. def start(self):
  642. if self.mp is None:
  643. self.mp = multiprocessing.Manager()
  644. for monitor in self.registry.values():
  645. monitor.start()
  646. self.started = True
  647. def stop(self):
  648. for monitor in self.registry.values():
  649. monitor.queue.close()
  650. monitor.terminate()
  651. if self.mp:
  652. self.mp.shutdown()
  653. self.mp = None
  654. self.started = False
  655. def notify_listener(self, taskno, msg):
  656. for monitor in self.registry.values():
  657. monitor.queue.put(('notify_listener', taskno, msg))
  658. def add_listener(self, no):
  659. for monitor in self.registry.values():
  660. monitor.queue.put(('add_listener', no))
  661. # ######################### plac server ############################# #
  662. import asyncore
  663. import asynchat
  664. import socket
  665. class _AsynHandler(asynchat.async_chat):
  666. "asynchat handler starting a new interpreter loop for each connection"
  667. terminator = '\r\n' # the standard one for telnet
  668. prompt = 'i> '
  669. def __init__(self, socket, interpreter):
  670. asynchat.async_chat.__init__(self, socket)
  671. self.set_terminator(self.terminator)
  672. self.i = interpreter
  673. self.i.__enter__()
  674. self.data = []
  675. self.write(self.prompt)
  676. def write(self, data, *args):
  677. "Push a string back to the client"
  678. if args:
  679. data %= args
  680. if data.endswith('\n') and not data.endswith(self.terminator):
  681. data = data[:-1] + self.terminator # fix newlines
  682. self.push(data)
  683. def collect_incoming_data(self, data):
  684. "Collect one character at the time"
  685. self.data.append(data)
  686. def found_terminator(self):
  687. "Put in the queue the line received from the client"
  688. line = ''.join(self.data)
  689. self.log('Received line %r from %s' % (line, self.addr))
  690. if line == 'EOF':
  691. self.i.__exit__(None, None, None)
  692. self.handle_close()
  693. else:
  694. task = self.i.submit(line)
  695. task.run() # synchronous or not
  696. if task.etype: # manage exception
  697. error = '%s: %s\nReceived: %s' % (
  698. task.etype.__name__, task.exc, ' '.join(task.arglist))
  699. self.log_info(task.traceback + error) # on the server
  700. self.write(error + self.terminator) # back to the client
  701. else: # no exception
  702. self.write(task.str + self.terminator)
  703. self.data = []
  704. self.write(self.prompt)
  705. class _AsynServer(asyncore.dispatcher):
  706. "asyncore-based server spawning AsynHandlers"
  707. def __init__(self, interpreter, newhandler, port, listen=5):
  708. self.interpreter = interpreter
  709. self.newhandler = newhandler
  710. self.port = port
  711. asyncore.dispatcher.__init__(self)
  712. self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
  713. self.bind(('', port))
  714. self.listen(listen)
  715. def handle_accept(self):
  716. clientsock, clientaddr = self.accept()
  717. self.log('Connected from %s' % str(clientaddr))
  718. i = self.interpreter.__class__(self.interpreter.obj) # new interpreter
  719. self.newhandler(clientsock, i) # spawn a new handler
  720. # ########################## the Interpreter ############################ #
  721. class Interpreter(object):
  722. """
  723. A context manager with a .send method and a few utility methods:
  724. execute, test and doctest.
  725. """
  726. class Exit(Exception):
  727. pass
  728. def __init__(self, obj, commentchar='#', split=shlex.split):
  729. self.obj = obj
  730. try:
  731. self.name = obj.__module__
  732. except AttributeError:
  733. self.name = 'plac'
  734. self.commentchar = commentchar
  735. self.split = split
  736. self._set_commands(obj)
  737. self.tm = TaskManager(obj)
  738. self.man = self.tm.man
  739. self.parser = self.tm.parser
  740. if self.commands:
  741. self.parser.addsubcommands(
  742. self.tm.specialcommands, self.tm, title='special commands')
  743. if obj.mpcommands:
  744. self.parser.addsubcommands(
  745. obj.mpcommands, obj,
  746. title='commands run in external processes')
  747. if obj.thcommands:
  748. self.parser.addsubcommands(
  749. obj.thcommands, obj, title='threaded commands')
  750. self.parser.error = lambda msg: sys.exit(msg) # patch the parser
  751. self._interpreter = None
  752. def _set_commands(self, obj):
  753. "Make sure obj has the right command attributes as Python sets"
  754. for attrname in ('commands', 'mpcommands', 'thcommands'):
  755. setattr(self, attrname, set(getattr(self.__class__, attrname, [])))
  756. setattr(obj, attrname, set(getattr(obj, attrname, [])))
  757. self.commands = obj.commands
  758. self.mpcommands.update(obj.mpcommands)
  759. self.thcommands.update(obj.thcommands)
  760. if (obj.commands or obj.mpcommands or obj.thcommands) and \
  761. not hasattr(obj, 'help'): # add default help
  762. obj.help = default_help.__get__(obj, obj.__class__)
  763. self.commands.add('help')
  764. def __enter__(self):
  765. "Start the inner interpreter loop"
  766. self._interpreter = self._make_interpreter()
  767. self._interpreter.send(None)
  768. return self
  769. def __exit__(self, exctype, exc, tb):
  770. "Close the inner interpreter and the task manager"
  771. self.close(exctype, exc, tb)
  772. def submit(self, line):
  773. "Send a line to the underlying interpreter and return a task object"
  774. if self._interpreter is None:
  775. raise RuntimeError(_('%r not initialized: probably you forgot to '
  776. 'use the with statement') % self)
  777. if isinstance(line, (str, bytes)):
  778. arglist = self.split(line, self.commentchar)
  779. else: # expects a list of strings
  780. arglist = line
  781. if not arglist:
  782. return nulltask
  783. m = self.tm.man # manager
  784. if m and not m.started:
  785. m.start()
  786. task = self._interpreter.send(arglist) # nonblocking
  787. if not plac_core._match_cmd(arglist[0], self.tm.specialcommands):
  788. self.tm.registry[task.no] = task
  789. if m:
  790. m.add_listener(task.no)
  791. return task
  792. def send(self, line):
  793. """Send a line to the underlying interpreter and return
  794. the finished task"""
  795. task = self.submit(line)
  796. BaseTask.run(task) # blocking
  797. return task
  798. def tasks(self):
  799. "The full lists of the submitted tasks"
  800. return self.tm.registry.values()
  801. def close(self, exctype=None, exc=None, tb=None):
  802. "Can be called to close the interpreter prematurely"
  803. self.tm.close()
  804. if exctype is not None:
  805. self._interpreter.throw(exctype, exc, tb)
  806. else:
  807. self._interpreter.close()
  808. def _make_interpreter(self):
  809. "The interpreter main loop, from lists of arguments to task objects"
  810. enter = getattr(self.obj, '__enter__', lambda: None)
  811. exit = getattr(self.obj, '__exit__', lambda et, ex, tb: None)
  812. enter()
  813. task = None
  814. try:
  815. for no in itertools.count(1):
  816. arglist = yield task
  817. try:
  818. cmd, result = self.parser.consume(arglist)
  819. except SystemExit as e: # for invalid commands
  820. if e.args == (0,): # raised as sys.exit(0)
  821. errlist = []
  822. else:
  823. errlist = [str(e)]
  824. task = SynTask(no, arglist, iter(errlist))
  825. continue
  826. except: # anything else
  827. task = SynTask(no, arglist, gen_exc(*sys.exc_info()))
  828. continue
  829. if not plac_core.iterable(result): # atomic result
  830. task = SynTask(no, arglist, gen_val(result))
  831. elif cmd in self.obj.mpcommands:
  832. task = MPTask(no, arglist, result, self.tm.man)
  833. elif cmd in self.obj.thcommands:
  834. task = ThreadedTask(no, arglist, result)
  835. else: # blocking task
  836. task = SynTask(no, arglist, result)
  837. except GeneratorExit: # regular exit
  838. exit(None, None, None)
  839. except: # exceptional exit
  840. exit(*sys.exc_info())
  841. raise
  842. def check(self, given_input, expected_output):
  843. "Make sure you get the expected_output from the given_input"
  844. output = self.send(given_input).str # blocking
  845. ok = (output == expected_output)
  846. if not ok:
  847. # the message here is not internationalized on purpose
  848. msg = 'input: %s\noutput: %s\nexpected: %s' % (
  849. given_input, output, expected_output)
  850. raise AssertionError(msg)
  851. def _parse_doctest(self, lineiter):
  852. "Returns the lines of input, the lines of output, and the line number"
  853. lines = [line.strip() for line in lineiter]
  854. inputs = []
  855. positions = []
  856. for i, line in enumerate(lines):
  857. if line.startswith('i> '):
  858. inputs.append(line[3:])
  859. positions.append(i)
  860. positions.append(len(lines) + 1) # last position
  861. outputs = []
  862. for i, start in enumerate(positions[:-1]):
  863. end = positions[i + 1]
  864. outputs.append('\n'.join(lines[start+1:end]))
  865. return zip(inputs, outputs, positions)
  866. def doctest(self, lineiter, verbose=False):
  867. """
  868. Parse a text containing doctests in a context and tests of all them.
  869. Raise an error even if a single doctest if broken. Use this for
  870. sequential tests which are logically grouped.
  871. """
  872. with self:
  873. try:
  874. for input, output, no in self._parse_doctest(lineiter):
  875. if verbose:
  876. write('i> %s\n' % input)
  877. write('-> %s\n' % output)
  878. task = self.send(input) # blocking
  879. if not str(task) == output:
  880. msg = ('line %d: input: %s\noutput: %s\nexpected: %s\n'
  881. % (no + 1, input, task, output))
  882. write(msg)
  883. if task.exc:
  884. raise_(task.etype, task.exc, task.tb)
  885. except self.Exit:
  886. pass
  887. def execute(self, lineiter, verbose=False):
  888. "Execute a lineiter of commands in a context and print the output"
  889. with self:
  890. try:
  891. for line in lineiter:
  892. if verbose:
  893. write('i> ' + line)
  894. task = self.send(line) # finished task
  895. if task.etype: # there was an error
  896. raise_(task.etype, task.exc, task.tb)
  897. write('%s\n' % task.str)
  898. except self.Exit:
  899. pass
  900. def multiline(self, stdin=sys.stdin, terminator=';', verbose=False):
  901. "The multiline mode is especially suited for usage with emacs"
  902. with self:
  903. try:
  904. for line in read_long_line(stdin, terminator):
  905. task = self.submit(line)
  906. task.run()
  907. write('%s\n' % task.str)
  908. if verbose and task.traceback:
  909. write(task.traceback)
  910. except self.Exit:
  911. pass
  912. def interact(self, stdin=sys.stdin, prompt='i> ', verbose=False):
  913. "Starts an interactive command loop reading commands from the consolle"
  914. try:
  915. import readline
  916. readline_present = True
  917. except ImportError:
  918. readline_present = False
  919. if stdin is sys.stdin and readline_present: # use readline
  920. histfile = os.path.expanduser('~/.%s.history' % self.name)
  921. completions = list(self.commands) + list(self.mpcommands) + \
  922. list(self.thcommands) + list(self.tm.specialcommands)
  923. self.stdin = ReadlineInput(completions, histfile=histfile)
  924. else:
  925. self.stdin = stdin
  926. self.prompt = prompt
  927. self.verbose = verbose
  928. intro = self.obj.__doc__ or ''
  929. write(intro + '\n')
  930. with self:
  931. self.obj._interact_ = True
  932. if self.stdin is sys.stdin: # do not close stdin automatically
  933. self._manage_input()
  934. else:
  935. with self.stdin: # close stdin automatically
  936. self._manage_input()
  937. def _manage_input(self):
  938. "Convert input lines into task which are then executed"
  939. try:
  940. for line in iter(lambda: read_line(self.stdin, self.prompt), ''):
  941. line = line.strip()
  942. if not line:
  943. continue
  944. task = self.submit(line)
  945. task.run() # synchronous or not
  946. write(str(task) + '\n')
  947. if self.verbose and task.etype:
  948. write(task.traceback)
  949. except self.Exit:
  950. pass
  951. def start_server(self, port=2199, **kw):
  952. """Starts an asyncore server reading commands for clients and opening
  953. a new interpreter for each connection."""
  954. _AsynServer(self, _AsynHandler, port) # register the server
  955. try:
  956. asyncore.loop(**kw)
  957. except (KeyboardInterrupt, TerminatedProcess):
  958. pass
  959. finally:
  960. asyncore.close_all()
  961. def add_monitor(self, mon):
  962. self.man.add(mon)
  963. def del_monitor(self, name):
  964. self.man.delete(name)
  965. @classmethod
  966. def call(cls, factory, arglist=sys.argv[1:],
  967. commentchar='#', split=shlex.split,
  968. stdin=sys.stdin, prompt='i> ', verbose=False):
  969. """
  970. Call a container factory with the arglist and instantiate an
  971. interpreter object. If there are remaining arguments, send them to the
  972. interpreter, else start an interactive session.
  973. """
  974. obj = partial_call(factory, arglist)
  975. i = cls(obj, commentchar, split)
  976. if i.obj._args_:
  977. with i:
  978. task = i.send(i.obj._args_) # synchronous
  979. if task.exc:
  980. raise_(task.etype, task.exc, task.tb)
  981. out = str(task)
  982. if out:
  983. print(out)
  984. elif i.obj._interact_:
  985. i.interact(stdin, prompt, verbose)
  986. else:
  987. i.parser.print_usage()
  988. # ################################## runp ################################### #
  989. class _TaskLauncher(object):
  990. "Helper for runp"
  991. def __init__(self, genseq, mode):
  992. if mode == 'p':
  993. self.mpcommands = ['rungen']
  994. else:
  995. self.thcommands = ['rungen']
  996. self.genlist = list(genseq)
  997. def rungen(self, i):
  998. for out in self.genlist[int(i) - 1]:
  999. yield out
  1000. def runp(genseq, mode='p'):
  1001. """Run a sequence of generators in parallel. Mode can be 'p' (use processes)
  1002. or 't' (use threads). After all of them are finished, return a list of
  1003. task objects.
  1004. """
  1005. assert mode in 'pt', mode
  1006. launcher = _TaskLauncher(genseq, mode)
  1007. res = []
  1008. with Interpreter(launcher) as inter:
  1009. for i in range(len(launcher.genlist)):
  1010. inter.submit('rungen %d' % (i + 1)).run()
  1011. for task in inter.tasks():
  1012. try:
  1013. res.append(task.result)
  1014. except Exception as e:
  1015. res.append(e)
  1016. return res