Added options to return or include raw data from query

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,27 +195,60 @@ 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):
             data = data.decode(encoding)
         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("<remark>(?P<msg>[^<>]*)</remark>").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"):
         """

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

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