Newer
Older
exporter / exporter.py
#! /usr/bin/python

import socket
import sys
import os
import os.path
import json
import traceback
import shlex
from args import *
from job import *
from driver import *

top_level = {("--port","-p"):"wrapper_communication_port", \
             ("--debug","-d"):None}

per_arg = {("--out","-o"):"out_dir", \
           ("--base","-b"):"base_dir"}

top_level_all = [k[0][2:] for k in top_level]

def desc_args(args):
  return " ".join(["[" + "/".join(k) + ("" if args[k] is None else " " + args[k]) + "]" for k in args.keys()])

def list_long_args(args):
  long_args = [k[0] for k in args]
  ret = ", ".join(long_args[:-1])
  if len(long_args) >= 2:
    ret += " or " + long_args[-1]
  return ret

top_level_desc = desc_args(top_level)
top_level_args = list_long_args(top_level)

per_arg_desc = desc_args(per_arg)
      
def usage():
  print "usage 1: ./exporter.py " + top_level_desc + " " + per_arg_desc + " input_file1.src,input_file2.src,..."
  print "  Exports the assets specified in the given source file(s)"
  print
  print "usage 2: ./exporter.py " + top_level_desc + " -"
  print "  Reads full command lines (everything except " + per_arg_desc + ") as lines from stdin"
  print
  print "usage 3: ./export.py  --help|-h"
  print "  Shows this help message"

if "-h" in sys.argv or "--help" in sys.argv:
  usage()
  sys.exit(0)

(err, args) = ArgParser().parse_args(sys.argv[1:])
if err is not None:
  print err
  usage()
  sys.exit(2)

from_stdin = False
if "-" in args["input"]:
  from_stdin = True
  if len(args["input"]) > 1:
    print "If reading from stdin ('-') is specified, no other input may be given"
    usage()
    sys.exit(2)

if not len(args["input"]):
  usage()
  sys.exit(0)

port = args["port"] if "port" in args else 7890

jobs = {}
sent_all_jobs = False
driver = Driver(port, from_stdin, exporter_already_running = "debug" in args)

def send_cmd(cmd, **kwargs):
  global driver
  kwargs["command"] = cmd
  driver.send_json(kwargs)

def send_job(job):
  send_cmd("job", **job.get_cmd())

def check_jobs_done():
  global sent_all_, jobs
  return sent_all_jobs and not len(jobs)

if not from_stdin:
  input = args.pop("input")
  for i in input:
    try:
      job = Job(in_file = i, **args)
      jobs[job.id] = job
    except Exception, e:
      print e
  sent_all_jobs = True
  if check_jobs_done():
    sys.exit(0)

driver_started = driver.start()

if driver_started:
  if not from_stdin:
    for id in jobs:
      send_job(jobs[id])

def handle_stdin(line):
  global sent_all_jobs, jobs
  if line is None:
    sent_all_jobs = True
    if check_jobs_done():
      return False
    return True
  parts = shlex.split(line)
  (err, line_args) = ArgParser().parse_args(parts)
  if err is not None:
    print err
  else:
    for a in top_level_all:
      if a in line_args:
        print "Top level args may not be passed in stdin: " + top_level_args
        return True
    input = line_args.pop("input")

    if "out_dir" in args and not "out_dir" in line_args:
      line_args["out_dir"] = args["out_dir"]
    if "base_dir" in args and not "base_dir" in line_args:
      line_args["base_dir"] = args["base_dir"]

    for i in input:
      try:
        job = Job(in_file = i, **line_args)
        jobs[job.id] = job
        send_job(job)
      except Exception, e:
        print e
  return True

def handle_command(msg_str):
  # return False to quit the driver loop
  msg = json.loads(msg_str)
  if "command" in msg:
    cmd = msg["command"]
    if cmd == "exit":
      return False
    if cmd == "print" and "string" in msg:
      sys.stdout.write(msg["string"])
      sys.stdout.flush()
    if cmd == "done" and "id" in msg:
      id = int(msg["id"])
      if id in jobs:
        job = jobs[id]
        jobs.pop(id)

        job.write_results({} if "data" not in msg else msg["data"], [] if "resources" not in msg else msg["resources"])

        if check_jobs_done():
          return False
  return True

# still have to loop() even if we didn't start succesfully,
# as we may be waiting for the exporter process to close
driver.loop(handle_command, handle_stdin)