Skip to content

Instantly share code, notes, and snippets.

@RussellLuo

RussellLuo/grpc_pi.py

Last active April 15, 2020 10:35
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save RussellLuo/9ee9585e3c2b0dbd0298574c241e1bcf to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save RussellLuo/9ee9585e3c2b0dbd0298574c241e1bcf to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. RussellLuo revised this gist May 11, 2017. 1 changed file with 30 additions and 12 deletions.
    42 changes: 30 additions & 12 deletions grpc_pi.py
    View file
    Original file line number Diff line number Diff line change
    @@ -21,10 +21,11 @@ class Generator(object):
    writer = sys.stdout

    def __init__(self, proto_package_name, pb2_module_name,
    core_method_name, rpc_method_args_size):
    core_method_name, unfold_method_args, rpc_method_args_size):
    self.proto_package_name = proto_package_name
    self.pb2_module_name = pb2_module_name
    self.core_method_name = core_method_name
    self.unfold_method_args = unfold_method_args
    self.rpc_method_args_size = rpc_method_args_size

    self.stub_class_name = self.camelize(self.proto_package_name) + 'Stub'
    @@ -145,9 +146,7 @@ def write_stub_property(self):

    def write_core_method(self):
    self.writer.write(
    '\n def {core_method_name}(self, rpc_name, req_name, **kwargs):\n'
    ' req_class = getattr({pb2_name}, req_name)\n'
    ' req = req_class(**kwargs)\n\n'
    '\n def {core_method_name}(self, rpc_name, req):\n'
    ' rpc = getattr(self.stub, rpc_name)\n'
    ' resp = rpc(req, self.timeout)\n'
    ' return resp\n'.format(
    @@ -156,7 +155,19 @@ def write_core_method(self):
    )
    )

    def write_rpc_method(self, method_name, req_name, req_param_names):
    def write_folded_rpc_method(self, method_name, req_name):
    self.writer.write(
    "\n def {underscored_method_name}(self, {req_name}):\n"
    " resp = self.{core_method_name}('{method_name}', {req_name})\n"
    " return resp\n".format(
    underscored_method_name=self.underscore(method_name),
    req_name=self.underscore(req_name),
    core_method_name=self.core_method_name,
    method_name=method_name
    )
    )

    def write_unfolded_rpc_method(self, method_name, req_name, req_param_names):
    indented_header = ' def {}('.format(self.underscore(method_name))

    full_params = ['self'] + req_param_names
    @@ -172,16 +183,15 @@ def write_rpc_method(self, method_name, req_name, req_param_names):
    for param_name in req_param_names
    )
    indented_body = (
    " resp = self.{core_method_name}(\n"
    " '{method_name}',\n"
    " '{req_name}',\n"
    " req = {req_name}(\n"
    "{indented_kwargs}\n"
    " )\n"
    " resp = self.{core_method_name}('{method_name}', req)\n"
    " return resp\n".format(
    core_method_name=self.core_method_name,
    method_name=method_name,
    req_name=req_name,
    indented_kwargs=indented_kwargs
    indented_kwargs=indented_kwargs,
    core_method_name=self.core_method_name,
    method_name=method_name
    )
    )
    self.writer.write(
    @@ -217,7 +227,11 @@ def write_rpc_methods(self):
    self.underscore(field.name)
    for field in req_class.DESCRIPTOR.fields
    ]
    self.write_rpc_method(stub_method_name, req_name, req_param_names)
    if self.unfold_method_args:
    self.write_unfolded_rpc_method(stub_method_name, req_name,
    req_param_names)
    else:
    self.write_folded_rpc_method(stub_method_name, req_name)

    def generate(self):
    self.write_module_header()
    @@ -240,13 +254,17 @@ def main():
    parser.add_argument('--core-method-name', default='call_rpc',
    help='The name of the core method that will be '
    'used to call the actual rpc methods.')
    parser.add_argument('--unfold-method-args', action='store_true',
    help='Whether or not to unfold the request '
    'attributes as the arguments of each rpc method.')
    parser.add_argument('--rpc-method-args-size', type=int, default=0,
    help='The number of arguments per line in the '
    'definition of each rpc method.')
    args = parser.parse_args()
    generator = Generator(args.proto_package_name,
    args.pb2_module_name,
    args.core_method_name,
    args.unfold_method_args,
    args.rpc_method_args_size)
    generator.generate()

  2. RussellLuo revised this gist May 11, 2017. 1 changed file with 2 additions and 2 deletions.
    4 changes: 2 additions & 2 deletions grpc_pi.py
    View file
    Original file line number Diff line number Diff line change
    @@ -86,11 +86,11 @@ def write_module_header(self):
    else:
    import_pb2 = 'import {pb2_name}'.format(pb2_name=self.pb2_name)
    self.writer.write(
    '# -*- coding: utf-8 -*-'
    '# -*- coding: utf-8 -*-\n'
    '{import_enum}'
    '\nimport grpc'
    '\n\n{import_pb2}'.format(
    import_enum='\n\nimport enum' if self.has_enum_types() else '',
    import_enum='\nimport enum' if self.has_enum_types() else '',
    import_pb2=import_pb2
    )
    )
  3. RussellLuo revised this gist May 11, 2017. 1 changed file with 3 additions and 2 deletions.
    5 changes: 3 additions & 2 deletions grpc_pi.py
    View file
    Original file line number Diff line number Diff line change
    @@ -88,8 +88,9 @@ def write_module_header(self):
    self.writer.write(
    '# -*- coding: utf-8 -*-'
    '{import_enum}'
    '\nimport grpc'
    '\n\n{import_pb2}'.format(
    import_enum='\n\nfrom enum import Enum' if self.has_enum_types() else '',
    import_enum='\n\nimport enum' if self.has_enum_types() else '',
    import_pb2=import_pb2
    )
    )
    @@ -103,7 +104,7 @@ def write_enum_types(self):
    for value in enum.values
    )
    self.writer.write(
    '\n\n\nclass {enum_name}(Enum):\n'
    '\n\n\nclass {enum_name}(enum.Enum):\n'
    '{values}'.format(enum_name=enum.name, values=values)
    )

  4. RussellLuo revised this gist May 11, 2017. 1 changed file with 69 additions and 18 deletions.
    87 changes: 69 additions & 18 deletions grpc_pi.py
    View file
    Original file line number Diff line number Diff line change
    @@ -3,7 +3,7 @@
    """Generate a pythonic interface based on the code generated by `grpcio-tools`.
    Example:
    $ python grpc_pi.py --pb2-module-name='python.path.xx_pb2' --stub-class-name='XxStub'
    $ python grpc_pi.py --proto-package-name='xx' --pb2-module-name='python.path.xx_pb2'
    """

    import argparse
    @@ -20,18 +20,23 @@ class Generator(object):

    writer = sys.stdout

    def __init__(self, pb2_module_name, stub_class_name,
    def __init__(self, proto_package_name, pb2_module_name,
    core_method_name, rpc_method_args_size):
    self.proto_package_name = proto_package_name
    self.pb2_module_name = pb2_module_name
    self.stub_class_name = stub_class_name
    self.core_method_name = core_method_name
    self.rpc_method_args_size = rpc_method_args_size

    self.stub_class_name = self.camelize(self.proto_package_name) + 'Stub'

    if '.' in self.pb2_module_name:
    self.pb2_path, self.pb2_name = self.pb2_module_name.rsplit('.', 1)
    else:
    self.pb2_path, self.pb2_name = '', self.pb2_module_name

    self.pb2_module = import_module(self.pb2_module_name)
    self.sym_db_pool = self.pb2_module._sym_db.pool

    @staticmethod
    def slice_every(iterable, n, padding=False, padding_item=None):
    """Return a list with at most `n` items each time from the `iterable`."""
    @@ -45,40 +50,85 @@ def slice_every(iterable, n, padding=False, padding_item=None):
    piece.extend([padding_item] * padding_len)
    yield piece

    @staticmethod
    def camelize(string, uppercase_first_letter=True):
    """Convert strings to CamelCase.
    Borrowed from https://github.com/jpvanhal/inflection/blob/master/inflection.py
    """
    if uppercase_first_letter:
    return re.sub(r"(?:^|_)(.)", lambda m: m.group(1).upper(), string)
    else:
    return string[0].lower() + Generator.camelize(string)[1:]

    @staticmethod
    def underscore(word):
    """Make an underscored, lowercase form from the expression
    in the string.
    Borrowed from https://github.com/jpvanhal/inflection/blob/master/inflection.py
    """
    word = re.sub(r"([A-Z]+)([A-Z][a-z])", r'\1_\2', word)
    word = re.sub(r"([a-z\d])([A-Z])", r'\1_\2', word)
    word = word.replace("-", "_")
    return word.lower()

    def has_enum_types(self):
    return any(name.startswith(self.proto_package_name)
    for name in self.sym_db_pool._enum_descriptors)

    def write_module_header(self):
    if self.pb2_path:
    import_stmt = 'from {pb2_path} import {pb2_name}'.format(
    import_pb2 = 'from {pb2_path} import {pb2_name}'.format(
    pb2_path=self.pb2_path,
    pb2_name=self.pb2_name
    )
    else:
    import_stmt = 'import {pb2_name}'.format(pb2_name=self.pb2_name)
    import_pb2 = 'import {pb2_name}'.format(pb2_name=self.pb2_name)
    self.writer.write(
    '# -*- coding=utf-8 -*-'
    '\n\n{import_stmt}'.format(import_stmt=import_stmt)
    '# -*- coding: utf-8 -*-'
    '{import_enum}'
    '\n\n{import_pb2}'.format(
    import_enum='\n\nfrom enum import Enum' if self.has_enum_types() else '',
    import_pb2=import_pb2
    )
    )

    def write_enum_types(self):
    for name, enum in self.sym_db_pool._enum_descriptors.iteritems():
    if name.startswith(self.proto_package_name):
    values = '\n'.join(
    ' {name} = {number}'.format(name=value.name,
    number=value.number)
    for value in enum.values
    )
    self.writer.write(
    '\n\n\nclass {enum_name}(Enum):\n'
    '{values}'.format(enum_name=enum.name, values=values)
    )

    def write_message_types(self):
    self.writer.write('\n\n')
    for name, message in self.sym_db_pool._descriptors.iteritems():
    if name.startswith(self.proto_package_name):
    self.writer.write(
    '\n{name} = {pb2_name}.{name}'.format(
    name=message.name,
    pb2_name=self.pb2_name
    )
    )

    def write_class_header(self):
    class_prefix = self.stub_class_name.rstrip('Stub')
    class_prefix = self.camelize(self.proto_package_name)
    self.writer.write(
    '\n\n\nclass {}Interface(object):\n'
    '\n timeout = 10\n'.format(class_prefix)
    '\n\n\nclass {}Interface(object):\n'.format(class_prefix)
    )

    def write_class_constructor(self):
    self.writer.write(
    '\n def __init__(self, target):'
    '\n self.target = target\n'
    '\n def __init__(self, target, timeout=10):'
    '\n self.target = target'
    '\n self.timeout = timeout\n'
    )

    def write_stub_property(self):
    @@ -144,8 +194,7 @@ def write_rpc_method(self, method_name, req_name, req_param_names):
    )

    def write_rpc_methods(self):
    pb2_module = import_module(self.pb2_module_name)
    stub_class = getattr(pb2_module, self.stub_class_name)
    stub_class = getattr(self.pb2_module, self.stub_class_name)

    channel = grpc.insecure_channel('localhost')
    stub = stub_class(channel)
    @@ -171,6 +220,8 @@ def write_rpc_methods(self):

    def generate(self):
    self.write_module_header()
    self.write_enum_types()
    self.write_message_types()
    self.write_class_header()
    self.write_class_constructor()
    self.write_stub_property()
    @@ -180,20 +231,20 @@ def generate(self):

    def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--proto-package-name', required=True,
    help='The package name of the proto file.')
    parser.add_argument('--pb2-module-name', required=True,
    help='The name of the generated `xx_pdb2.py` '
    'module with the full Python path.')
    parser.add_argument('--stub-class-name', required=True,
    help='The name of the generated `XxStub` class.')
    parser.add_argument('--core-method-name', default='call_rpc',
    help='The name of the core method that will be '
    'used to call the actual rpc methods.')
    parser.add_argument('--rpc-method-args-size', type=int, default=0,
    help='The number of arguments per line in the '
    'definition of each rpc method.')
    args = parser.parse_args()
    generator = Generator(args.pb2_module_name,
    args.stub_class_name,
    generator = Generator(args.proto_package_name,
    args.pb2_module_name,
    args.core_method_name,
    args.rpc_method_args_size)
    generator.generate()
  5. RussellLuo revised this gist May 11, 2017. 1 changed file with 9 additions and 2 deletions.
    11 changes: 9 additions & 2 deletions grpc_pi.py
    View file
    Original file line number Diff line number Diff line change
    @@ -75,12 +75,18 @@ def write_class_header(self):
    '\n timeout = 10\n'.format(class_prefix)
    )

    def write_class_constructor(self):
    self.writer.write(
    '\n def __init__(self, target):'
    '\n self.target = target\n'
    )

    def write_stub_property(self):
    self.writer.write(
    '\n @property\n'
    ' def stub(self):\n'
    ' # channel = grpc.insecure_channel(...)\n'
    ' # return {pb2_name}.{stub_class_name}(channel)\n'.format(
    ' channel = grpc.insecure_channel(self.target)\n'
    ' return {pb2_name}.{stub_class_name}(channel)\n'.format(
    pb2_name=self.pb2_name,
    stub_class_name=self.stub_class_name
    )
    @@ -166,6 +172,7 @@ def write_rpc_methods(self):
    def generate(self):
    self.write_module_header()
    self.write_class_header()
    self.write_class_constructor()
    self.write_stub_property()
    self.write_core_method()
    self.write_rpc_methods()
  6. RussellLuo revised this gist Sep 3, 2016. 1 changed file with 101 additions and 0 deletions.
    101 changes: 101 additions & 0 deletions mocker.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,101 @@
    # -*- coding=utf-8 -*-

    """A mocking gRPC client interface for tests.
    Usage:
    1. Suppose you want to test the RPC function `create_user` of gRPC
    client interface `UserInterface`, then the test may look like:
    ```
    import unittest
    from user_module import UserInterface
    class TestUserInterface(unittest.TestCase):
    user_inter = UserInterface(...)
    def test_create_user(self):
    result = self.user_inter.create_user(name='russellluo')
    self.assertEqual(result, 0)
    ```
    The above test will pass if the real gRPC server is running, but
    running a gRPC server for tests is cumbersome.
    2. Use `Mocker` to simply the test environment
    ```
    import unittest
    from user_module import UserInterface
    from mocker_module import Mocker
    class TestUserInterface(unittest.TestCase):
    user_inter = Mocker(UserInterface(...))
    def test_create_user(self):
    result = self.user_inter.create_user(name='russellluo')
    self.assertEqual(result, 0)
    ```
    1) First, run the test once, by interacting with the real gRPC server
    2) Then, you will find that the class `Mocker` in the module file
    `mocker_module` is changed magically
    3) Afterwards, you can run the test alone as many times as you wish,
    and the test will always pass without any interaction with the real
    gRPC server
    Internals:
    When interacting with the real gRPC server, the mocking class
    `Mocker` can automatically record the real input/output data of
    each call of each RPC function, and finally `Mocker` will change
    itself to be a complete replacement for the real gRPC interface.
    """

    import cPickle
    import functools
    import os
    import pprint
    from collections import defaultdict


    def record(data, method):
    @functools.wraps(method)
    def decorator(*args, **kwargs):
    params = cPickle.dumps(args) + cPickle.dumps(kwargs)
    method_data = data[method.__name__]
    if params not in method_data:
    result = method(*args, **kwargs)
    method_data[params] = cPickle.dumps(result)
    return cPickle.loads(method_data[params])
    return decorator


    class Mocker(object):

    def __init__(self, target):
    self._data = getattr(self, '_fake_data', defaultdict(dict))
    for attr_name in dir(target):
    if not attr_name.startswith('_'):
    target_method = getattr(target, attr_name)
    recordable_func = record(self._data, target_method)
    setattr(self, attr_name, recordable_func)

    def __del__(self):
    mocker_file = os.path.abspath(__file__)
    with open(mocker_file, 'r') as f:
    content = f.read()

    with open(mocker_file, 'w') as f:
    fake_data_comment = '# Generated fake data'
    mocker_body = ''.join(content.rpartition(fake_data_comment)[:-1])
    f.write(mocker_body)

    fake_data_string = pprint.pformat(dict(self._data))
    f.write('\n _fake_data = ' + fake_data_string)

    # Generated fake data
  7. RussellLuo renamed this gist Aug 31, 2016. 1 changed file with 0 additions and 0 deletions.
    0 grpc_mpi.py → grpc_pi.py
    View file
    File renamed without changes.
  8. RussellLuo renamed this gist Aug 31, 2016. 1 changed file with 2 additions and 1 deletion.
    3 changes: 2 additions & 1 deletion grpc_pi.py → grpc_mpi.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,6 +1,6 @@
    # -*- coding=utf-8 -*-

    """Generate a more pythonic interface based on the code generated by `grpcio-tools`.
    """Generate a pythonic interface based on the code generated by `grpcio-tools`.
    Example:
    $ python grpc_pi.py --pb2-module-name='python.path.xx_pb2' --stub-class-name='XxStub'
    @@ -34,6 +34,7 @@ def __init__(self, pb2_module_name, stub_class_name,

    @staticmethod
    def slice_every(iterable, n, padding=False, padding_item=None):
    """Return a list with at most `n` items each time from the `iterable`."""
    iterable = iter(iterable)
    while True:
    piece = list(itertools.islice(iterable, n))
  9. RussellLuo revised this gist Aug 31, 2016. 1 changed file with 46 additions and 15 deletions.
    61 changes: 46 additions & 15 deletions grpc_pi.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,13 +1,13 @@
    # -*- coding=utf-8 -*-

    """Generate a more pythonic interface based on the code generated by
    `grpcio-tools`.
    """Generate a more pythonic interface based on the code generated by `grpcio-tools`.
    Example:
    $ python grpc_pi.py --pb2-module-name='python.path.xx_pb2' --stub-class-name='XxStub'
    """

    import argparse
    import itertools
    import re
    import sys
    from collections import OrderedDict
    @@ -20,12 +20,29 @@ class Generator(object):

    writer = sys.stdout

    def __init__(self, pb2_module_name, stub_class_name, core_method_name):
    def __init__(self, pb2_module_name, stub_class_name,
    core_method_name, rpc_method_args_size):
    self.pb2_module_name = pb2_module_name
    self.stub_class_name = stub_class_name
    self.core_method_name = core_method_name
    self.rpc_method_args_size = rpc_method_args_size

    self.pb2_path, self.pb2_module = self.pb2_module_name.rsplit('.', 1)
    if '.' in self.pb2_module_name:
    self.pb2_path, self.pb2_name = self.pb2_module_name.rsplit('.', 1)
    else:
    self.pb2_path, self.pb2_name = '', self.pb2_module_name

    @staticmethod
    def slice_every(iterable, n, padding=False, padding_item=None):
    iterable = iter(iterable)
    while True:
    piece = list(itertools.islice(iterable, n))
    if not piece:
    return
    padding_len = n - len(piece)
    if padding_len and padding:
    piece.extend([padding_item] * padding_len)
    yield piece

    @staticmethod
    def underscore(word):
    @@ -38,9 +55,16 @@ def underscore(word):
    return word.lower()

    def write_module_header(self):
    if self.pb2_path:
    import_stmt = 'from {pb2_path} import {pb2_name}'.format(
    pb2_path=self.pb2_path,
    pb2_name=self.pb2_name
    )
    else:
    import_stmt = 'import {pb2_name}'.format(pb2_name=self.pb2_name)
    self.writer.write(
    '# -*- coding=utf-8 -*-'
    '\n\nfrom {} import {}'.format(self.pb2_path, self.pb2_module)
    '\n\n{import_stmt}'.format(import_stmt=import_stmt)
    )

    def write_class_header(self):
    @@ -55,31 +79,34 @@ def write_stub_property(self):
    '\n @property\n'
    ' def stub(self):\n'
    ' # channel = grpc.insecure_channel(...)\n'
    ' # return {pb2_module}.{stub_class_name}(channel)\n'.format(
    pb2_module=self.pb2_module,
    ' # return {pb2_name}.{stub_class_name}(channel)\n'.format(
    pb2_name=self.pb2_name,
    stub_class_name=self.stub_class_name
    )
    )

    def write_core_method(self):
    self.writer.write(
    '\n def {core_method_name}(self, rpc_name, req_name, **kwargs):\n'
    ' req_class = getattr({pb2_module}, req_name)\n'
    ' req_class = getattr({pb2_name}, req_name)\n'
    ' req = req_class(**kwargs)\n\n'
    ' rpc = getattr(self.stub, rpc_name)\n'
    ' resp = rpc(req, self.timeout)\n'
    ' return resp\n'.format(
    core_method_name=self.core_method_name,
    pb2_module=self.pb2_module
    pb2_name=self.pb2_name
    )
    )

    def write_rpc_method(self, method_name, req_name, req_param_names):
    indented_header = ' def {}('.format(self.underscore(method_name))
    indent_len = len(indented_header)
    indented_params = ',\n'.join(
    (indent_len * ' ') + param_name
    for param_name in req_param_names

    full_params = ['self'] + req_param_names
    args_size = self.rpc_method_args_size or len(full_params)
    separator = ',\n' + len(indented_header) * ' '
    indented_params = separator.join(
    ', '.join(params)
    for params in self.slice_every(full_params, args_size)
    )

    indented_kwargs = ',\n'.join(
    @@ -100,7 +127,7 @@ def write_rpc_method(self, method_name, req_name, req_param_names):
    )
    )
    self.writer.write(
    '\n{indented_header}self,\n'
    '\n{indented_header}'
    '{indented_params}):\n'
    '{indented_body}'.format(
    indented_header=indented_header,
    @@ -153,10 +180,14 @@ def main():
    parser.add_argument('--core-method-name', default='call_rpc',
    help='The name of the core method that will be '
    'used to call the actual rpc methods.')
    parser.add_argument('--rpc-method-args-size', type=int, default=0,
    help='The number of arguments per line in the '
    'definition of each rpc method.')
    args = parser.parse_args()
    generator = Generator(args.pb2_module_name,
    args.stub_class_name,
    args.core_method_name)
    args.core_method_name,
    args.rpc_method_args_size)
    generator.generate()


  10. RussellLuo revised this gist Aug 30, 2016. 1 changed file with 2 additions and 3 deletions.
    5 changes: 2 additions & 3 deletions grpc_pi.py
    View file
    Original file line number Diff line number Diff line change
    @@ -75,7 +75,7 @@ def write_core_method(self):
    )

    def write_rpc_method(self, method_name, req_name, req_param_names):
    indented_header = ' def {}('.format(method_name)
    indented_header = ' def {}('.format(self.underscore(method_name))
    indent_len = len(indented_header)
    indented_params = ',\n'.join(
    (indent_len * ' ') + param_name
    @@ -133,8 +133,7 @@ def write_rpc_methods(self):
    self.underscore(field.name)
    for field in req_class.DESCRIPTOR.fields
    ]
    method_name = self.underscore(stub_method_name)
    self.write_rpc_method(method_name, req_name, req_param_names)
    self.write_rpc_method(stub_method_name, req_name, req_param_names)

    def generate(self):
    self.write_module_header()
  11. RussellLuo revised this gist Aug 30, 2016. 1 changed file with 160 additions and 1 deletion.
    161 changes: 160 additions & 1 deletion grpc_pi.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,5 +1,164 @@
    # -*- coding=utf-8 -*-

    """Generate a more pythonic interface based on the code generated by
    `grpcio-tools`.
    Example:
    $ python grpc_pi.py --pb2-module-name='python.path.xx_pb2' --stub-class-name='XxStub'
    """

    import argparse
    import re
    import sys
    from collections import OrderedDict
    from importlib import import_module

    import grpc


    class Generator(object):

    writer = sys.stdout

    def __init__(self, pb2_module_name, stub_class_name, core_method_name):
    self.pb2_module_name = pb2_module_name
    self.stub_class_name = stub_class_name
    self.core_method_name = core_method_name

    self.pb2_path, self.pb2_module = self.pb2_module_name.rsplit('.', 1)

    @staticmethod
    def underscore(word):
    """Make an underscored, lowercase form from the expression
    in the string.
    """
    word = re.sub(r"([A-Z]+)([A-Z][a-z])", r'\1_\2', word)
    word = re.sub(r"([a-z\d])([A-Z])", r'\1_\2', word)
    word = word.replace("-", "_")
    return word.lower()

    def write_module_header(self):
    self.writer.write(
    '# -*- coding=utf-8 -*-'
    '\n\nfrom {} import {}'.format(self.pb2_path, self.pb2_module)
    )

    def write_class_header(self):
    class_prefix = self.stub_class_name.rstrip('Stub')
    self.writer.write(
    '\n\n\nclass {}Interface(object):\n'
    '\n timeout = 10\n'.format(class_prefix)
    )

    def write_stub_property(self):
    self.writer.write(
    '\n @property\n'
    ' def stub(self):\n'
    ' # channel = grpc.insecure_channel(...)\n'
    ' # return {pb2_module}.{stub_class_name}(channel)\n'.format(
    pb2_module=self.pb2_module,
    stub_class_name=self.stub_class_name
    )
    )

    def write_core_method(self):
    self.writer.write(
    '\n def {core_method_name}(self, rpc_name, req_name, **kwargs):\n'
    ' req_class = getattr({pb2_module}, req_name)\n'
    ' req = req_class(**kwargs)\n\n'
    ' rpc = getattr(self.stub, rpc_name)\n'
    ' resp = rpc(req, self.timeout)\n'
    ' return resp\n'.format(
    core_method_name=self.core_method_name,
    pb2_module=self.pb2_module
    )
    )

    def write_rpc_method(self, method_name, req_name, req_param_names):
    indented_header = ' def {}('.format(method_name)
    indent_len = len(indented_header)
    indented_params = ',\n'.join(
    (indent_len * ' ') + param_name
    for param_name in req_param_names
    )

    indented_kwargs = ',\n'.join(
    ' {0}={0}'.format(param_name)
    for param_name in req_param_names
    )
    indented_body = (
    " resp = self.{core_method_name}(\n"
    " '{method_name}',\n"
    " '{req_name}',\n"
    "{indented_kwargs}\n"
    " )\n"
    " return resp\n".format(
    core_method_name=self.core_method_name,
    method_name=method_name,
    req_name=req_name,
    indented_kwargs=indented_kwargs
    )
    )
    self.writer.write(
    '\n{indented_header}self,\n'
    '{indented_params}):\n'
    '{indented_body}'.format(
    indented_header=indented_header,
    indented_params=indented_params,
    indented_body=indented_body
    )
    )

    def write_rpc_methods(self):
    pb2_module = import_module(self.pb2_module_name)
    stub_class = getattr(pb2_module, self.stub_class_name)

    channel = grpc.insecure_channel('localhost')
    stub = stub_class(channel)
    stub_method_names = [
    attr
    for attr in dir(stub)
    if not attr.startswith('__')
    ]
    stub_method_names.sort()
    stub_methods = OrderedDict([
    (stub_method_name, getattr(stub, stub_method_name))
    for stub_method_name in stub_method_names
    ])

    for stub_method_name, stub_method in stub_methods.iteritems():
    req_class = stub_method._request_serializer.im_class
    req_name = req_class.__name__
    req_param_names = [
    self.underscore(field.name)
    for field in req_class.DESCRIPTOR.fields
    ]
    method_name = self.underscore(stub_method_name)
    self.write_rpc_method(method_name, req_name, req_param_names)

    def generate(self):
    self.write_module_header()
    self.write_class_header()
    self.write_stub_property()
    self.write_core_method()
    self.write_rpc_methods()


    def main():
    pass
    parser = argparse.ArgumentParser()
    parser.add_argument('--pb2-module-name', required=True,
    help='The name of the generated `xx_pdb2.py` '
    'module with the full Python path.')
    parser.add_argument('--stub-class-name', required=True,
    help='The name of the generated `XxStub` class.')
    parser.add_argument('--core-method-name', default='call_rpc',
    help='The name of the core method that will be '
    'used to call the actual rpc methods.')
    args = parser.parse_args()
    generator = Generator(args.pb2_module_name,
    args.stub_class_name,
    args.core_method_name)
    generator.generate()


    if __name__ == '__main__':
  12. RussellLuo revised this gist Aug 30, 2016. No changes.
  13. RussellLuo renamed this gist Aug 30, 2016. 1 changed file with 0 additions and 0 deletions.
    0 grpc_proto_to_class.py → grpc_pi.py
    View file
    File renamed without changes.
  14. RussellLuo created this gist Aug 29, 2016.
    6 changes: 6 additions & 0 deletions grpc_proto_to_class.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,6 @@
    def main():
    pass


    if __name__ == '__main__':
    main()