diff --git a/overpy/__init__.py b/overpy/__init__.py index a85bb01..3c6d744 100644 --- a/overpy/__init__.py +++ b/overpy/__init__.py @@ -111,12 +111,12 @@ def _handle_remark_msg(msg: str) -> NoReturn: raise exception.OverpassRuntimeRemark(msg=msg) raise exception.OverpassUnknownError(msg=msg) - def query(self, query: Union[bytes, str]) -> "Result": + def query_raw(self, query: Union[bytes, str]) -> Tuple[Union[bytes, str], str]: """ - Query the Overpass API + Query the Overpass API and returns the raw response :param query: The query string in Overpass QL - :return: The parsed result + :return: A tuple made of the raw response and the response content type """ if not isinstance(query, bytes): query = query.encode("utf-8") @@ -145,11 +145,8 @@ def query(self, query: Union[bytes, str]) -> "Result": if f.code == 200: content_type = f.getheader("Content-Type") - if content_type == "application/json": - return self.parse_json(response) - - if content_type == "application/osm3s+xml": - return self.parse_xml(response) + if content_type in ("application/json", "application/osm3s+xml"): + return response, content_type current_exception = exception.OverpassUnknownContentType(content_type) if not do_retry: @@ -198,12 +195,36 @@ def query(self, query: Union[bytes, str]) -> "Result": raise exception.MaxRetriesReached(retry_count=retry_num, exceptions=retry_exceptions) - def parse_json(self, data: Union[bytes, str], encoding: str = "utf-8") -> "Result": + def query( + self, + query: Union[bytes, str], + *, + include_raw: bool = False) -> "Result": + """ + Query the Overpass API and returns parsed result + + :param query: The query string in Overpass QL + :param include_raw: True to store the raw data along with the parsed result + :return: The parsed result + """ + + response, content_type = self.query_raw(query) + if content_type == "application/json": + return self.parse_json(response, include_raw=include_raw) + else: # "application/osm3s+xml" + return self.parse_xml(response, include_raw=include_raw) + + def parse_json( + self, + data: Union[bytes, str], + encoding: str = "utf-8", + include_raw: bool = False) -> "Result": """ Parse raw response from Overpass service. :param data: Raw JSON Data :param encoding: Encoding to decode byte string + :param include_raw: True to store the data along with the parsed result :return: Result object """ if isinstance(data, bytes): @@ -211,14 +232,23 @@ def parse_json(self, data: Union[bytes, str], encoding: str = "utf-8") -> "Resul data_parsed: dict = json.loads(data, parse_float=Decimal) if "remark" in data_parsed: self._handle_remark_msg(msg=data_parsed.get("remark")) - return Result.from_json(data_parsed, api=self) + result = Result.from_json(data_parsed, api=self) + if include_raw: + result.raw = data + return result - def parse_xml(self, data: Union[bytes, str], encoding: str = "utf-8", parser: Optional[int] = None): + def parse_xml( + self, + data: Union[bytes, str], + encoding: str = "utf-8", + parser: Optional[int] = None, + include_raw: bool = False) -> "Result": """ :param data: Raw XML Data :param encoding: Encoding to decode byte string :param parser: The XML parser to use + :param include_raw: True to store the data along with the parsed result :return: Result object """ if parser is None: @@ -230,8 +260,10 @@ def parse_xml(self, data: Union[bytes, str], encoding: str = "utf-8", parser: Op m = re.compile("(?P[^<>]*)").search(data) if m: self._handle_remark_msg(m.group("msg")) - - return Result.from_xml(data, api=self, parser=parser) + result = Result.from_xml(data, api=self, parser=parser) + if include_raw: + result.raw = data + return result class Result: @@ -242,11 +274,13 @@ class Result: def __init__( self, elements: Optional[List[Union["Area", "Node", "Relation", "Way"]]] = None, - api: Optional[Overpass] = None): + api: Optional[Overpass] = None, + raw: Optional[Union[bytes, str]] = None): """ :param elements: List of elements to initialize the result with :param api: The API object to load additional resources and elements + :param raw: The raw data corresponding to these elements """ if elements is None: elements = [] @@ -269,6 +303,7 @@ def __init__( Area: self._areas } self.api = api + self.raw = raw def expand(self, other: "Result"): """ diff --git a/tests/test_json.py b/tests/test_json.py index 9a96d0e..97f35a2 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -116,3 +116,17 @@ def test_remark_unknown(self): api = overpy.Overpass() with pytest.raises(overpy.exception.OverpassUnknownError): api.parse_json(read_file("json/remark-unknown-01.json")) + + +class TestIncludeRawJSON: + def test_include_raw_json(self): + api = overpy.Overpass() + data = read_file("json/area-01.json") + result = api.parse_json(data, include_raw=True) + assert result.raw == data + + def test_not_include_raw(self): + api = overpy.Overpass() + data = read_file("json/area-01.json") + result = api.parse_json(data, include_raw=False) + assert result.raw is None diff --git a/tests/test_xml.py b/tests/test_xml.py index 8a00b76..a3db45d 100644 --- a/tests/test_xml.py +++ b/tests/test_xml.py @@ -209,3 +209,17 @@ def test_remark_unknown(self): api = overpy.Overpass() with pytest.raises(overpy.exception.OverpassUnknownError): api.parse_xml(read_file("xml/remark-unknown-01.xml")) + + +class TestIncludeRawXML: + def test_include_raw_xml(self): + api = overpy.Overpass() + data = read_file("xml/area-01.xml") + result = api.parse_xml(data, include_raw=True) + assert result.raw == data + + def test_not_include_raw(self): + api = overpy.Overpass() + data = read_file("xml/area-01.xml") + result = api.parse_xml(data, include_raw=False) + assert result.raw is None