void t_py_generator::generate_service_client()

in compiler/cpp/src/thrift/generate/t_py_generator.cc [1388:1717]


void t_py_generator::generate_service_client(t_service* tservice) {
  string extends = "";
  string extends_client = "";
  if (tservice->get_extends() != nullptr) {
    extends = type_name(tservice->get_extends());
    if (gen_zope_interface_) {
      extends_client = "(" + extends + ".Client)";
    } else {
      extends_client = extends + ".Client, ";
    }
  } else {
    if (gen_zope_interface_ && (gen_newstyle_ || gen_dynamic_)) {
      extends_client = "(object)";
    }
  }

  f_service_ << '\n' << '\n';

  if (gen_zope_interface_) {
    f_service_ << "@implementer(Iface)" << '\n'
               << "class Client" << extends_client << ":" << '\n'
               << '\n';
  } else {
    f_service_ << "class Client(" << extends_client << "Iface):" << '\n';
  }
  indent_up();
  generate_python_docstring(f_service_, tservice);

  // Constructor function
  if (gen_twisted_) {
    f_service_ << indent() << "def __init__(self, transport, oprot_factory):" << '\n';
  } else if (gen_tornado_) {
    f_service_ << indent()
               << "def __init__(self, transport, iprot_factory, oprot_factory=None):" << '\n';
  } else {
    f_service_ << indent() << "def __init__(self, iprot, oprot=None):" << '\n';
  }
  indent_up();
  if (extends.empty()) {
    if (gen_twisted_) {
      f_service_ << indent() << "self._transport = transport" << '\n'
                 << indent() << "self._oprot_factory = oprot_factory" << '\n'
                 << indent() << "self._seqid = 0" << '\n'
                 << indent() << "self._reqs = {}" << '\n';
    } else if (gen_tornado_) {
      f_service_ << indent() << "self._transport = transport" << '\n'
                 << indent() << "self._iprot_factory = iprot_factory" << '\n'
                 << indent() << "self._oprot_factory = (oprot_factory if oprot_factory is not None"
                 << '\n'
                 << indent() << "                       else iprot_factory)" << '\n'
                 << indent() << "self._seqid = 0" << '\n'
                 << indent() << "self._reqs = {}" << '\n'
                 << indent() << "self._transport.io_loop.spawn_callback(self._start_receiving)"
                 << '\n';
    } else {
      f_service_ << indent() << "self._iprot = self._oprot = iprot" << '\n'
                 << indent() << "if oprot is not None:" << '\n'
                 << indent() << indent_str() << "self._oprot = oprot" << '\n'
                 << indent() << "self._seqid = 0" << '\n';
    }
  } else {
    if (gen_twisted_) {
      f_service_ << indent() << extends
                 << ".Client.__init__(self, transport, oprot_factory)" << '\n';
    } else if (gen_tornado_) {
      f_service_ << indent() << extends
                 << ".Client.__init__(self, transport, iprot_factory, oprot_factory)" << '\n';
    } else {
      f_service_ << indent() << extends << ".Client.__init__(self, iprot, oprot)" << '\n';
    }
  }
  indent_down();

  if (gen_tornado_ && extends.empty()) {
    f_service_ << '\n' <<
      indent() << "@gen.coroutine" << '\n' <<
      indent() << "def _start_receiving(self):" << '\n';
    indent_up();
    indent(f_service_) << "while True:" << '\n';
    indent_up();
    f_service_ << indent() << "try:" << '\n'
               << indent() << indent_str() << "frame = yield self._transport.readFrame()" << '\n'
               << indent() << "except TTransport.TTransportException as e:" << '\n'
               << indent() << indent_str() << "for future in self._reqs.values():" << '\n'
               << indent() << indent_str() << indent_str() << "future.set_exception(e)" << '\n'
               << indent() << indent_str() << "self._reqs = {}" << '\n'
               << indent() << indent_str() << "return" << '\n'
               << indent() << "tr = TTransport.TMemoryBuffer(frame)" << '\n'
               << indent() << "iprot = self._iprot_factory.getProtocol(tr)" << '\n'
               << indent() << "(fname, mtype, rseqid) = iprot.readMessageBegin()" << '\n'
               << indent() << "method = getattr(self, 'recv_' + fname)" << '\n'
               << indent() << "future = self._reqs.pop(rseqid, None)" << '\n'
               << indent() << "if not future:" << '\n'
               << indent() << indent_str() << "# future has already been discarded" << '\n'
               << indent() << indent_str() << "continue" << '\n'
               << indent() << "try:" << '\n'
               << indent() << indent_str() << "result = method(iprot, mtype, rseqid)" << '\n'
               << indent() << "except Exception as e:" << '\n'
               << indent() << indent_str() << "future.set_exception(e)" << '\n'
               << indent() << "else:" << '\n'
               << indent() << indent_str() << "future.set_result(result)" << '\n';
    indent_down();
    indent_down();
  }

  // Generate client method implementations
  vector<t_function*> functions = tservice->get_functions();
  vector<t_function*>::const_iterator f_iter;
  for (f_iter = functions.begin(); f_iter != functions.end(); ++f_iter) {
    t_struct* arg_struct = (*f_iter)->get_arglist();
    const vector<t_field*>& fields = arg_struct->get_members();
    vector<t_field*>::const_iterator fld_iter;
    string funname = (*f_iter)->get_name();

    f_service_ << '\n';
    // Open function
    indent(f_service_) << "def " << function_signature(*f_iter, false) << ":" << '\n';
    indent_up();
    generate_python_docstring(f_service_, (*f_iter));
    if (gen_twisted_) {
      indent(f_service_) << "seqid = self._seqid = self._seqid + 1" << '\n';
      indent(f_service_) << "self._reqs[seqid] = defer.Deferred()" << '\n' << '\n';
      indent(f_service_) << "d = defer.maybeDeferred(self.send_" << funname;

    } else if (gen_tornado_) {
      indent(f_service_) << "self._seqid += 1" << '\n';
      if (!(*f_iter)->is_oneway()) {
        indent(f_service_) << "future = self._reqs[self._seqid] = concurrent.Future()" << '\n';
      }
      indent(f_service_) << "self.send_" << funname << "(";

    } else {
      indent(f_service_) << "self.send_" << funname << "(";
    }

    bool first = true;
    if (gen_twisted_) {
      // we need a leading comma if there are args, since it's called as maybeDeferred(funcname,
      // arg)
      first = false;
    }
    for (fld_iter = fields.begin(); fld_iter != fields.end(); ++fld_iter) {
      if (first) {
        first = false;
      } else {
        f_service_ << ", ";
      }
      f_service_ << (*fld_iter)->get_name();
    }

    f_service_ << ")" << '\n';

    if (!(*f_iter)->is_oneway()) {
      if (gen_twisted_) {
        // nothing. See the next block.
      } else if (gen_tornado_) {
        indent(f_service_) << "return future" << '\n';
      } else {
        f_service_ << indent();
        if (!(*f_iter)->get_returntype()->is_void()) {
          f_service_ << "return ";
        }
        f_service_ << "self.recv_" << funname << "()" << '\n';
      }
    }
    indent_down();

    if (gen_twisted_) {
      // This block injects the body of the send_<> method for twisted (and a cb/eb pair)
      indent_up();
      indent(f_service_) << "d.addCallbacks(" << '\n';

      indent_up();
      f_service_ << indent() << "callback=self.cb_send_" << funname << "," << '\n' << indent()
                 << "callbackArgs=(seqid,)," << '\n' << indent() << "errback=self.eb_send_"
                 << funname << "," << '\n' << indent() << "errbackArgs=(seqid,))" << '\n';
      indent_down();

      indent(f_service_) << "return d" << '\n';
      indent_down();
      f_service_ << '\n';

      indent(f_service_) << "def cb_send_" << funname << "(self, _, seqid):" << '\n';
      indent_up();
      if ((*f_iter)->is_oneway()) {
        // if one-way, fire the deferred & remove it from _reqs
        f_service_ << indent() << "d = self._reqs.pop(seqid)" << '\n' << indent()
                   << "d.callback(None)" << '\n' << indent() << "return d" << '\n';
      } else {
        f_service_ << indent() << "return self._reqs[seqid]" << '\n';
      }
      indent_down();
      f_service_ << '\n';

      // add an errback to fail the request if the call to send_<> raised an exception
      indent(f_service_) << "def eb_send_" << funname << "(self, f, seqid):" << '\n';
      indent_up();
      f_service_ << indent() << "d = self._reqs.pop(seqid)" << '\n' << indent() << "d.errback(f)"
                 << '\n' << indent() << "return d" << '\n';
      indent_down();
    }

    f_service_ << '\n';
    indent(f_service_) << "def send_" << function_signature(*f_iter, false, true) << ":" << '\n';
    indent_up();

    std::string argsname = (*f_iter)->get_name() + "_args";
    std::string messageType = (*f_iter)->is_oneway() ? "TMessageType.ONEWAY" : "TMessageType.CALL";

    // Serialize the request header
    if (gen_twisted_ || gen_tornado_) {
      f_service_ << indent() << "oprot = self._oprot_factory.getProtocol(self._transport)" << '\n'
                 << indent() << "oprot.writeMessageBegin('" << (*f_iter)->get_name() << "', "
                 << messageType << ", self._seqid)" << '\n';
    } else {
      f_service_ << indent() << "self._oprot.writeMessageBegin('" << (*f_iter)->get_name() << "', "
                 << messageType << ", self._seqid)" << '\n';
    }

    f_service_ << indent() << "args = " << argsname << "()" << '\n';

    for (fld_iter = fields.begin(); fld_iter != fields.end(); ++fld_iter) {
      f_service_ << indent() << "args." << (*fld_iter)->get_name() << " = "
                 << (*fld_iter)->get_name() << '\n';
    }

    // Write to the stream
    if (gen_twisted_ || gen_tornado_) {
      f_service_ << indent() << "args.write(oprot)" << '\n' << indent() << "oprot.writeMessageEnd()"
                 << '\n' << indent() << "oprot.trans.flush()" << '\n';
    } else {
      f_service_ << indent() << "args.write(self._oprot)" << '\n' << indent()
                 << "self._oprot.writeMessageEnd()" << '\n' << indent()
                 << "self._oprot.trans.flush()" << '\n';
    }

    indent_down();

    if (!(*f_iter)->is_oneway()) {
      std::string resultname = (*f_iter)->get_name() + "_result";
      // Open function
      f_service_ << '\n';
      if (gen_twisted_ || gen_tornado_) {
        f_service_ << indent() << "def recv_" << (*f_iter)->get_name()
                   << "(self, iprot, mtype, rseqid):" << '\n';
      } else {
        t_struct noargs(program_);
        t_function recv_function((*f_iter)->get_returntype(),
                                 string("recv_") + (*f_iter)->get_name(),
                                 &noargs);
        f_service_ << indent() << "def " << function_signature(&recv_function) << ":" << '\n';
      }
      indent_up();

      // TODO(mcslee): Validate message reply here, seq ids etc.

      if (gen_twisted_) {
        f_service_ << indent() << "d = self._reqs.pop(rseqid)" << '\n';
      } else if (gen_tornado_) {
      } else {
        f_service_ << indent() << "iprot = self._iprot" << '\n' << indent()
                   << "(fname, mtype, rseqid) = iprot.readMessageBegin()" << '\n';
      }

      f_service_ << indent() << "if mtype == TMessageType.EXCEPTION:" << '\n'
                 << indent() << indent_str() << "x = TApplicationException()" << '\n';

      if (gen_twisted_) {
        f_service_ << indent() << indent_str() << "x.read(iprot)" << '\n' << indent()
                   << indent_str() << "iprot.readMessageEnd()" << '\n' << indent() << indent_str() << "return d.errback(x)"
                   << '\n' << indent() << "result = " << resultname << "()" << '\n' << indent()
                   << "result.read(iprot)" << '\n' << indent() << "iprot.readMessageEnd()" << '\n';
      } else {
        f_service_ << indent() << indent_str() << "x.read(iprot)" << '\n' << indent()
                   << indent_str() << "iprot.readMessageEnd()" << '\n' << indent() << indent_str() << "raise x" << '\n'
                   << indent() << "result = " << resultname << "()" << '\n' << indent()
                   << "result.read(iprot)" << '\n' << indent() << "iprot.readMessageEnd()" << '\n';
      }

      // Careful, only return _result if not a void function
      if (!(*f_iter)->get_returntype()->is_void()) {
        f_service_ << indent() << "if result.success is not None:" << '\n';
        if (gen_twisted_) {
          f_service_ << indent() << indent_str() << "return d.callback(result.success)" << '\n';
        } else {
          f_service_ << indent() << indent_str() << "return result.success" << '\n';
        }
      }

      t_struct* xs = (*f_iter)->get_xceptions();
      const std::vector<t_field*>& xceptions = xs->get_members();
      vector<t_field*>::const_iterator x_iter;
      for (x_iter = xceptions.begin(); x_iter != xceptions.end(); ++x_iter) {
        const string& xname = (*x_iter)->get_name();
        f_service_ << indent() << "if result." << xname << " is not None:" << '\n';
        if (gen_twisted_) {
          f_service_ << indent() << indent_str() << "return d.errback(result." << xname << ")"
                     << '\n';
        } else {
          f_service_ << indent() << indent_str() << "raise result." << xname << "" << '\n';
        }
      }

      // Careful, only return _result if not a void function
      if ((*f_iter)->get_returntype()->is_void()) {
        if (gen_twisted_) {
          f_service_ << indent() << "return d.callback(None)" << '\n';
        } else {
          f_service_ << indent() << "return" << '\n';
        }
      } else {
        if (gen_twisted_) {
          f_service_
              << indent()
              << "return d.errback(TApplicationException(TApplicationException.MISSING_RESULT, \""
              << (*f_iter)->get_name() << " failed: unknown result\"))" << '\n';
        } else {
          f_service_ << indent()
                     << "raise TApplicationException(TApplicationException.MISSING_RESULT, \""
                     << (*f_iter)->get_name() << " failed: unknown result\")" << '\n';
        }
      }

      // Close function
      indent_down();
    }
  }

  indent_down();
}