Merge pull request #26 from hhatto/provide-enum-for-input-output-type

Provide enum for input output type

README.rst

@@ -126,6 +126,9 @@ and call with ``output_type="apache-arrow"`` option:
         output_columns="_key,name",
     )
 
+.. note:: `output_type` can also specify `poyonga.const.OutputType.APACHE_ARROW`.
+    If is not specified, `poyonga.const.OutputType.JSON` is specified by default.
+
 load with ``input_type="apache-arrow"``:
 
 .. code-block:: python
@@ -145,6 +148,9 @@ load with ``input_type="apache-arrow"``:
     g = Groonga()
     g.call("load", table="Site", values=values, input_type="apache-arrow")
 
+.. note:: `input_type` can also specify `poyonga.const.InputType.APACHE_ARROW`.
+    If is not specified, `poyonga.const.InputType.JSON` is specified by default.
+
 
 more information:
 

examples/groonga_microblog_tutorial/2_load.py

@@ -1,3 +1,4 @@
+from poyonga import InputType
 from poyonga.client import Groonga
 
 
@@ -46,7 +47,7 @@ def _call(g, cmd, **kwargs):
 """
 
 
-_call(g, "load", table="Users", values="".join(users.splitlines()))
+_call(g, "load", table="Users", values="".join(users.splitlines()), input_type=InputType.JSON)
 _call(g, "select", table="Users")
 
 

examples/groonga_microblog_tutorial/3_search_error_with_pyarrow.py

@@ -1,5 +1,5 @@
-# coding: utf-8
 from poyonga import Groonga
+from poyonga.const import OutputType
 
 
 def _call_with_apachearrow(g, cmd, **kwargs):
@@ -27,7 +27,7 @@ def _call_with_apachearrow(g, cmd, **kwargs):
     table="Comment",  # NOTE: invalid table name
     filter="last_modified<=1268802000",
     output_columns="posted_by.name,comment,last_modified",
-    output_type="apache-arrow",
+    output_type=OutputType.APACHE_ARROW,
     drilldown="hash_tags,posted_by",
     drilldown_output_column="_id",
 )

poyonga/__init__.py

@@ -1,6 +1,7 @@
 from poyonga.client import Groonga
+from poyonga.const import InputType, OutputType
 from poyonga.result import GroongaResult
 
 __version__ = "0.5.0"
 
-__all__ = ["Groonga", "GroongaResult"]
+__all__ = ["Groonga", "GroongaResult", "InputType", "OutputType"]

poyonga/client.py

@@ -11,7 +11,7 @@
 except ImportError:
     import json
 
-from poyonga.const import GQTP_HEADER_SIZE
+from poyonga.const import GQTP_HEADER_SIZE, InputType, OutputType
 from poyonga.result import GroongaResult, GroongaSelectResult
 
 
@@ -113,7 +113,7 @@ def _call_http(self, cmd, **kwargs):
                 post_data = kwargs.pop("values")
             url = "".join([url, "?", urlencode(kwargs)])
         if post_data:
-            if kwargs.get("input_type") == "apache-arrow":
+            if kwargs.get("input_type") == InputType.APACHE_ARROW:
                 content_type = "application/x-apache-arrow-streaming"
             else:
                 content_type = "application/json"
@@ -137,7 +137,7 @@ def _call_http(self, cmd, **kwargs):
     def call(self, cmd, **kwargs):
         output_type = kwargs.get("output_type")
         if not output_type:
-            output_type = "json"
+            output_type = OutputType.JSON
         if self.protocol == "http" or self.protocol == "https":
             metadata, data = self._call_http(cmd, **kwargs)
         else:

poyonga/const.py

@@ -1,3 +1,11 @@
+try:
+    from enum import StrEnum
+except ImportError:
+    from enum import Enum
+    class StrEnum(str, Enum):  # noqa
+        pass
+
+
 # protocol type
 GRN_PROTO_GQTP = "gqtp"
 GRN_PROTO_HTTP = "http"
@@ -79,3 +87,15 @@
 GRN_STATUS_TOO_SMALL_LIMIT = 65467
 GRN_STATUS_CAS_ERROR = 65466
 GRN_STATUS_UNSUPPORTED_COMMAND_VERSION = 65465
+
+
+class InputType(StrEnum):
+    JSON = "json"
+    APACHE_ARROW = "apache-arrow"
+
+
+class OutputType(StrEnum):
+    JSON = "json"
+    TSV = "tsv"
+    MSGPACK = "msgpack"
+    APACHE_ARROW = "apache-arrow"

poyonga/result.py

@@ -20,25 +20,27 @@
 except ImportError:
     pa = None
 
+from poyonga.const import OutputType
+
 
 class GroongaResult:
-    def __init__(self, data, output_type="json", encoding="utf-8", content_type=None):
+    def __init__(self, data, output_type: OutputType = OutputType.JSON, encoding="utf-8", content_type=None):
         self.raw_result = data
-        if output_type == "tsv" or content_type == "text/tab-separated-values":
+        if output_type == OutputType.TSV or content_type == "text/tab-separated-values":
             # TODO: not implement
             csv.reader(StringIO(data), delimiter="\t")
             raise NotImplementedError(f"not implement output_type: {output_type}")
-        elif output_type == "msgpack" or content_type == "application/x-msgpack":
+        elif output_type == OutputType.MSGPACK or content_type == "application/x-msgpack":
             if msgpack:
                 _result = msgpack.unpackb(data)
             else:
                 raise Exception("msgpack is not supported")
-        elif output_type == "json" or content_type == "application/json":
+        elif output_type == OutputType.JSON or content_type == "application/json":
             if encoding == "utf-8":
                 _result = json.loads(data)
             else:
                 _result = json.loads(data, encoding=encoding)
-        elif output_type == "apache-arrow":
+        elif output_type == OutputType.APACHE_ARROW:
             if self._is_apache_arrow(content_type):
                 self._parse_apache_arrow(data)
                 return
@@ -105,7 +107,7 @@ def __str__(self):
 
 
 class GroongaSelectResult(GroongaResult):
-    def __init__(self, data, output_type="json", encoding="utf-8", content_type=None):
+    def __init__(self, data, output_type: OutputType = OutputType.JSON, encoding="utf-8", content_type=None):
         super(GroongaSelectResult, self).__init__(data, output_type, encoding, content_type)
         if self._is_apache_arrow(content_type):
             return