diff --git a/README.rst b/README.rst index 7c15641..4ac15dd 100644 --- a/README.rst +++ b/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: diff --git a/examples/groonga_microblog_tutorial/2_load.py b/examples/groonga_microblog_tutorial/2_load.py index d9c55af..bb07193 100644 --- a/examples/groonga_microblog_tutorial/2_load.py +++ b/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") diff --git a/examples/groonga_microblog_tutorial/3_search_error_with_pyarrow.py b/examples/groonga_microblog_tutorial/3_search_error_with_pyarrow.py index d635479..94ff01f 100644 --- a/examples/groonga_microblog_tutorial/3_search_error_with_pyarrow.py +++ b/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", ) diff --git a/poyonga/__init__.py b/poyonga/__init__.py index 0195ed4..af7df36 100644 --- a/poyonga/__init__.py +++ b/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"] diff --git a/poyonga/client.py b/poyonga/client.py index d8f2c30..fb7c64a 100644 --- a/poyonga/client.py +++ b/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: diff --git a/poyonga/const.py b/poyonga/const.py index 9db70bd..1dbb90f 100644 --- a/poyonga/const.py +++ b/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" diff --git a/poyonga/result.py b/poyonga/result.py index 3087285..12f75ed 100644 --- a/poyonga/result.py +++ b/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