Merge pull request #516 from c9s/narumi/fix-import-error

Fix: fix import error

Makefile

@@ -251,10 +251,12 @@ install-grpc-tools:
 	go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@v1.1
 	pip install grpcio-tools
 
+# https://github.com/protocolbuffers/protobuf/issues/1491#issuecomment-261914766
+# replace `import bbgo_pb2` by `from . import bbgo_pb2` to use relative import
 grpc-py:
 	python -m grpc_tools.protoc -I$(PWD)/pkg/pb \
-		--python_out=$(PWD)/python/bbgo \
+		--python_out=$(PWD)/python \
-		--grpc_python_out=$(PWD)/python/bbgo \
+		--grpc_python_out=$(PWD)/python \
 		$(PWD)/pkg/pb/bbgo.proto
 
 clean:

python/bbgo/__init__.py

@@ -1,2 +1,4 @@
-from .bbgo import BBGO
+from .services import MarketService
+from .services import TradingService
+from .services import UserDataService
 from .stream import Stream

python/bbgo/services.py

@@ -1,13 +1,48 @@
 from typing import List
 
-from . import bbgo_pb2
+import bbgo_pb2
-from .utils import create_stub
+import bbgo_pb2_grpc
 
 
-class BBGO(object):
+class UserDataService(object):
 
-    def __init__(self, host: str, port: int):
+    def __init__(self, stub: bbgo_pb2_grpc.UserDataServiceStub):
-        self.stub = create_stub(host, port)
+        self.stub = stub
+
+    def subscribe_user_data(self):
+        return self.stub.SubscribeUserData(bbgo_pb2.Empty())
+
+
+class MarketService(object):
+
+    def __init__(self, stub: bbgo_pb2_grpc.MarketDataServiceStub):
+        self.stub = stub
+
+    def subscribe(self, subscriptions: List[bbgo_pb2.Subscription]):
+        request = bbgo_pb2.SubscribeRequest(subscriptions=subscriptions)
+        request_iter = self.stub.Subscribe(request)
+        return request_iter
+
+    def query_klines(self,
+                     exchange: str,
+                     symbol: str,
+                     limit: int = 30,
+                     interval: int = 1,
+                     timestamp: int = None) -> bbgo_pb2.QueryKLinesResponse:
+        request = bbgo_pb2.QueryKLinesRequest(exchange=exchange,
+                                              symbol=symbol,
+                                              limit=limit,
+                                              interval=interval,
+                                              timestamp=timestamp)
+
+        response = self.market_data_stub.QueryKLines(request)
+        return response
+
+
+class TradingService(object):
+
+    def __init__(self, stub: bbgo_pb2_grpc.TradingServiceStub):
+        self.stub = stub
 
     def submit_order(self,
                      exchange: str,
@@ -87,18 +122,3 @@ class BBGO(object):
                                               offset=offset)
         response = self.stub.QueryTrades(request)
         return response
-
-    def query_klines(self,
-                     exchange: str,
-                     symbol: str,
-                     limit: int = 30,
-                     interval: int = 1,
-                     timestamp: int = None) -> bbgo_pb2.QueryKLinesResponse:
-        request = bbgo_pb2.QueryKLinesRequest(exchange=exchange,
-                                              symbol=symbol,
-                                              limit=limit,
-                                              interval=interval,
-                                              timestamp=timestamp)
-
-        response = self.stub.QueryKLines(request)
-        return response

python/bbgo/stream.py

@@ -2,11 +2,10 @@ import asyncio
 from typing import Callable
 from typing import List
 
+import bbgo_pb2
+import bbgo_pb2_grpc
 import grpc
 
-from . import bbgo_pb2
-from . import bbgo_pb2_grpc
-
 
 class Stream(object):
 
@@ -34,18 +33,18 @@ class Stream(object):
 
     async def subscribe(self):
         async with grpc.aio.insecure_channel(self.address) as channel:
-            stub = bbgo_pb2_grpc.BBGOStub(channel)
+            stub = bbgo_pb2_grpc.MarketDataServiceStub(channel)
 
             request = bbgo_pb2.SubscribeRequest(subscriptions=self.subscriptions)
-            async for response in stub.Subcribe(request):
+            async for response in stub.Subscribe(request):
                 self.dispatch(response)
 
     async def subscribe_user_data(self):
         async with grpc.aio.insecure_channel(self.address) as channel:
-            stub = bbgo_pb2_grpc.BBGOStub(channel)
+            stub = bbgo_pb2_grpc.UserDataServiceStub(channel)
 
             request = bbgo_pb2.Empty()
-            async for response in stub.SubcribeUserData(request):
+            async for response in stub.SubscribeUserData(request):
                 self.dispatch_user_events(response)
 
     def start(self):

python/bbgo/utils.py

@@ -2,8 +2,6 @@ import os
 
 import grpc
 
-from . import bbgo_pb2_grpc
-
 
 def read_binary(f):
     with open(f, 'rb') as fp:
@@ -29,9 +27,3 @@ def get_credentials_from_env():
     private_key_certificate_chain_pairs = [(private_key, certificate_chain)]
     server_credentials = grpc.ssl_server_credentials(private_key_certificate_chain_pairs)
     return server_credentials
-
-
-def create_stub(host, port):
-    address = f'{host}:{port}'
-    channel = grpc.insecure_channel(address)
-    return bbgo_pb2_grpc.BBGOStub(channel)

python/bbgo_pb2_grpc.py

@@ -15,12 +15,12 @@ class MarketDataServiceStub(object):
             channel: A grpc.Channel.
         """
         self.Subscribe = channel.unary_stream(
-                '/pb.MarketDataService/Subscribe',
+                '/bbgo.MarketDataService/Subscribe',
                 request_serializer=bbgo__pb2.SubscribeRequest.SerializeToString,
                 response_deserializer=bbgo__pb2.SubscribeResponse.FromString,
                 )
         self.QueryKLines = channel.unary_unary(
-                '/pb.MarketDataService/QueryKLines',
+                '/bbgo.MarketDataService/QueryKLines',
                 request_serializer=bbgo__pb2.QueryKLinesRequest.SerializeToString,
                 response_deserializer=bbgo__pb2.QueryKLinesResponse.FromString,
                 )
@@ -56,7 +56,7 @@ def add_MarketDataServiceServicer_to_server(servicer, server):
             ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
-            'pb.MarketDataService', rpc_method_handlers)
+            'bbgo.MarketDataService', rpc_method_handlers)
     server.add_generic_rpc_handlers((generic_handler,))
 
 
@@ -75,7 +75,7 @@ class MarketDataService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_stream(request, target, '/pb.MarketDataService/Subscribe',
+        return grpc.experimental.unary_stream(request, target, '/bbgo.MarketDataService/Subscribe',
             bbgo__pb2.SubscribeRequest.SerializeToString,
             bbgo__pb2.SubscribeResponse.FromString,
             options, channel_credentials,
@@ -92,7 +92,7 @@ class MarketDataService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/pb.MarketDataService/QueryKLines',
+        return grpc.experimental.unary_unary(request, target, '/bbgo.MarketDataService/QueryKLines',
             bbgo__pb2.QueryKLinesRequest.SerializeToString,
             bbgo__pb2.QueryKLinesResponse.FromString,
             options, channel_credentials,
@@ -109,7 +109,7 @@ class UserDataServiceStub(object):
             channel: A grpc.Channel.
         """
         self.SubscribeUserData = channel.unary_stream(
-                '/pb.UserDataService/SubscribeUserData',
+                '/bbgo.UserDataService/SubscribeUserData',
                 request_serializer=bbgo__pb2.Empty.SerializeToString,
                 response_deserializer=bbgo__pb2.SubscribeResponse.FromString,
                 )
@@ -135,7 +135,7 @@ def add_UserDataServiceServicer_to_server(servicer, server):
             ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
-            'pb.UserDataService', rpc_method_handlers)
+            'bbgo.UserDataService', rpc_method_handlers)
     server.add_generic_rpc_handlers((generic_handler,))
 
 
@@ -154,7 +154,7 @@ class UserDataService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_stream(request, target, '/pb.UserDataService/SubscribeUserData',
+        return grpc.experimental.unary_stream(request, target, '/bbgo.UserDataService/SubscribeUserData',
             bbgo__pb2.Empty.SerializeToString,
             bbgo__pb2.SubscribeResponse.FromString,
             options, channel_credentials,
@@ -171,27 +171,27 @@ class TradingServiceStub(object):
             channel: A grpc.Channel.
         """
         self.SubmitOrder = channel.unary_unary(
-                '/pb.TradingService/SubmitOrder',
+                '/bbgo.TradingService/SubmitOrder',
                 request_serializer=bbgo__pb2.SubmitOrderRequest.SerializeToString,
                 response_deserializer=bbgo__pb2.SubmitOrderResponse.FromString,
                 )
         self.CancelOrder = channel.unary_unary(
-                '/pb.TradingService/CancelOrder',
+                '/bbgo.TradingService/CancelOrder',
                 request_serializer=bbgo__pb2.CancelOrderRequest.SerializeToString,
                 response_deserializer=bbgo__pb2.CancelOrderResponse.FromString,
                 )
         self.QueryOrder = channel.unary_unary(
-                '/pb.TradingService/QueryOrder',
+                '/bbgo.TradingService/QueryOrder',
                 request_serializer=bbgo__pb2.QueryOrderRequest.SerializeToString,
                 response_deserializer=bbgo__pb2.QueryOrderResponse.FromString,
                 )
         self.QueryOrders = channel.unary_unary(
-                '/pb.TradingService/QueryOrders',
+                '/bbgo.TradingService/QueryOrders',
                 request_serializer=bbgo__pb2.QueryOrdersRequest.SerializeToString,
                 response_deserializer=bbgo__pb2.QueryOrdersResponse.FromString,
                 )
         self.QueryTrades = channel.unary_unary(
-                '/pb.TradingService/QueryTrades',
+                '/bbgo.TradingService/QueryTrades',
                 request_serializer=bbgo__pb2.QueryTradesRequest.SerializeToString,
                 response_deserializer=bbgo__pb2.QueryTradesResponse.FromString,
                 )
@@ -261,7 +261,7 @@ def add_TradingServiceServicer_to_server(servicer, server):
             ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
-            'pb.TradingService', rpc_method_handlers)
+            'bbgo.TradingService', rpc_method_handlers)
     server.add_generic_rpc_handlers((generic_handler,))
 
 
@@ -280,7 +280,7 @@ class TradingService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/pb.TradingService/SubmitOrder',
+        return grpc.experimental.unary_unary(request, target, '/bbgo.TradingService/SubmitOrder',
             bbgo__pb2.SubmitOrderRequest.SerializeToString,
             bbgo__pb2.SubmitOrderResponse.FromString,
             options, channel_credentials,
@@ -297,7 +297,7 @@ class TradingService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/pb.TradingService/CancelOrder',
+        return grpc.experimental.unary_unary(request, target, '/bbgo.TradingService/CancelOrder',
             bbgo__pb2.CancelOrderRequest.SerializeToString,
             bbgo__pb2.CancelOrderResponse.FromString,
             options, channel_credentials,
@@ -314,7 +314,7 @@ class TradingService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/pb.TradingService/QueryOrder',
+        return grpc.experimental.unary_unary(request, target, '/bbgo.TradingService/QueryOrder',
             bbgo__pb2.QueryOrderRequest.SerializeToString,
             bbgo__pb2.QueryOrderResponse.FromString,
             options, channel_credentials,
@@ -331,7 +331,7 @@ class TradingService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/pb.TradingService/QueryOrders',
+        return grpc.experimental.unary_unary(request, target, '/bbgo.TradingService/QueryOrders',
             bbgo__pb2.QueryOrdersRequest.SerializeToString,
             bbgo__pb2.QueryOrdersResponse.FromString,
             options, channel_credentials,
@@ -348,7 +348,7 @@ class TradingService(object):
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/pb.TradingService/QueryTrades',
+        return grpc.experimental.unary_unary(request, target, '/bbgo.TradingService/QueryTrades',
             bbgo__pb2.QueryTradesRequest.SerializeToString,
             bbgo__pb2.QueryTradesResponse.FromString,
             options, channel_credentials,

python/pyproject.toml

@@ -3,6 +3,11 @@ name = "bbgo"
 version = "0.1.0"
 description = ""
 authors = ["なるみ <weaper@gmail.com>"]
+packages = [
+    { include = "bbgo" },
+    { include = "bbgo_pb2.py" },
+    { include = "bbgo_pb2_grpc.py" },
+]
 
 [tool.poetry.dependencies]
 python = "^3.8"

python/tests/servicer.py

@@ -1,36 +1,8 @@
-import random
+import bbgo_pb2
-import time
+import bbgo_pb2_grpc
-
-from bbgo import bbgo_pb2
-from bbgo import bbgo_pb2_grpc
 
 
-class TestServicer(bbgo_pb2_grpc.BBGOServicer):
+class TestTradingServicer(bbgo_pb2_grpc.TradingServiceServicer):
-
-    def Subcribe(self, request, context):
-        i = 0
-        while True:
-            for subscription in request.subscriptions:
-                yield bbgo_pb2.SubscribeResponse(
-                    channel=subscription.channel,
-                    event=bbgo_pb2.Event.UPDATE,
-                    exchange=subscription.exchange,
-                    symbol=subscription.symbol + f'_{i}',
-                )
-            i += 1
-            time.sleep(random.random())
-
-    def SubcribeUserData(self, request, context):
-        i = 0
-        while True:
-            yield bbgo_pb2.SubscribeResponse(
-                channel=bbgo_pb2.Channel.USER,
-                event=bbgo_pb2.Event.ORDER_UPDATE,
-                exchange='max',
-                symbol=f'user_{i}',
-                )
-            i += 1
-            time.sleep(random.random())
 
     def SubmitOrder(self, request, context):
         submit_order = request.submit_order
@@ -50,18 +22,3 @@ class TestServicer(bbgo_pb2_grpc.BBGOServicer):
         error = bbgo_pb2.Error(error_code=0, error_message='')
 
         return bbgo_pb2.SubmitOrderResponse(order=order, error=error)
-
-    def CancelOrder(self, request, context):
-        pass
-
-    def QueryOrder(self, request, context):
-        pass
-
-    def QueryOrders(self, request, context):
-        pass
-
-    def QueryTrades(self, request, context):
-        pass
-
-    def QueryKLines(self, request, context):
-        pass

python/tests/test_grpc.py

@@ -1,48 +1,56 @@
 from concurrent import futures
 
+import bbgo_pb2
+import bbgo_pb2_grpc
 import grpc
 import pytest
 
-from bbgo import BBGO
+from bbgo import MarketService
-from bbgo import bbgo_pb2
+from bbgo import TradingService
-from bbgo import bbgo_pb2_grpc
+from tests.servicer import TestTradingServicer
-from tests.servicer import TestServicer
 
 
 @pytest.fixture
-def grpc_address(host='[::]', port=50051):
+def address(host='[::]', port=50051):
     return f'{host}:{port}'
 
 
 @pytest.fixture
-def bbgo(host='[::]', port=50051):
+def channel(address):
-    return BBGO(host, port)
+    return grpc.insecure_channel(address)
 
 
 @pytest.fixture
-def grpc_channel(grpc_address):
+def trading_service(channel):
-    return grpc.insecure_channel(grpc_address)
+    trading_service_stub = bbgo_pb2_grpc.TradingServiceStub(channel)
+    return TradingService(trading_service_stub)
 
 
 @pytest.fixture
-def grpc_server(grpc_address, max_workers=1):
+def market_service(channel):
+    market_service_stub = bbgo_pb2_grpc.MarketDataServiceStub(channel)
+    return MarketService(market_service_stub)
+
+
+@pytest.fixture
+def test_trading_servicer(address, max_workers=1):
     server = grpc.server(futures.ThreadPoolExecutor(max_workers))
-    servicer = TestServicer()
+    servicer = TestTradingServicer()
-    bbgo_pb2_grpc.add_BBGOServicer_to_server(servicer, server)
+    bbgo_pb2_grpc.add_TradingServiceServicer_to_server(servicer, server)
-    server.add_insecure_port(grpc_address)
+    server.add_insecure_port(address)
     server.start()
     yield server
     server.stop(grace=None)
 
 
-def test_submit_order(bbgo, grpc_server):
+def test_submit_order(trading_service, test_trading_servicer):
     exchange = 'max'
     symbol = 'BTCUSDT'
     side = bbgo_pb2.Side.BUY
     quantity = 0.01
     order_type = bbgo_pb2.OrderType.LIMIT
 
-    response = bbgo.submit_order(
+    response = trading_service.submit_order(
         exchange=exchange,
         symbol=symbol,
         side=side,