diff --git a/data_url/__init__.py b/data_url/__init__.py index 5b7df34..2a36e55 100644 --- a/data_url/__init__.py +++ b/data_url/__init__.py @@ -89,13 +89,17 @@ def from_data(cls, mime_type: str, base64_encoded: bool, data: str) -> 'DataURL' return data_url @classmethod - def from_byte_data(cls, mime_type: str, data: bytes) -> 'DataURL': - """Create a new data URL from a mime type and byte data. + def from_byte_data(cls, mime_type: str, data: Union[str, bytes]) -> 'DataURL': + """Create a new data URL from a mime type and raw data. - This method works similarly to from_data, however because the data is bytes type it will - automatically turn on base64 encoding. It also assumes that the data is not already - base64 encoded. If you have base64 encoded bytes convert them to a string then - use the `from_data` method. + If the data is a bytes type, then a data URL is constructed from the data's + base64 encoded representation. This assumes the data is not already base64 + encoded. If you have base64 encoded bytes, convert them to a string and use the + `from_data` method instead. + + If the data is a string type, then a data URL is constructed from the data's URL + encoded representation. This assumes the data is not already URL encoded. If you + have URL encoded data, use `from_data` instead. Args: mime_type (str) @@ -103,13 +107,17 @@ def from_byte_data(cls, mime_type: str, data: bytes) -> 'DataURL': Returns: DataURL: A new DataURL object. """ - if type(data) != bytes: - raise TypeError('Data must be a bytes type') - data_url = cls() data_url._mime_type = mime_type - data_url._is_base64_encoded = True - data_url._data = data + + if type(data) == str: + data_url._is_base64_encoded = False + data_url._data = quote(data) + elif type(data) == bytes: + data_url._is_base64_encoded = True + data_url._data = data + else: + raise TypeError('Data must be a string or bytes type') return data_url @@ -175,6 +183,13 @@ def encoded_data(self) -> str: return base64.b64encode(self._data).decode('utf-8') return self._data + @property + def decoded_data(self) -> Union[str, bytes]: + """The decoded data from the URL""" + if self._is_base64_encoded: + return self._data + return unquote(self._data) + @property def parameters(self) -> Dict[str, str]: """Attribute / Value parameters.""" diff --git a/test/test_url.py b/test/test_url.py index f03813e..95c0c8b 100644 --- a/test/test_url.py +++ b/test/test_url.py @@ -79,6 +79,12 @@ def test_url_assembly_with_parameters(self): url.parameters["name"] = "two words" self.assertEqual(str(url), self.example_url) + def test_url_with_url_escaped_encoding(self): + url = DataURL.from_url("data:,A%20brief%20note") + self.assertEqual("", url.mime_type) + self.assertEqual(False, url.is_base64_encoded) + self.assertEqual("A brief note", url.decoded_data) + class TestFromData(unittest.TestCase): def test_typing(self): @@ -175,16 +181,30 @@ def test_byte_data(self): self.assertEqual(type(self.url.data), bytes) self.run_assertions() + def test_string_with_spaces(self): + self.mime_type = "text/plain" + self.base64_encoded = False + self.data = "A brief note" + self.raw_data = quote(self.data) + self.expected_url = f"data:{self.mime_type},{self.raw_data}" + + self.url = DataURL.from_byte_data(self.mime_type, self.data) + self.assertEqual(type(self.url.data), str) + self.run_assertions() + def test_typing(self): with self.assertRaises(Exception) as context: - DataURL.from_byte_data("type", "string") + DataURL.from_byte_data("type", 1) - self.assertTrue('Data must be a bytes type' in str(context.exception)) + self.assertTrue('Data must be a string or bytes type' in str(context.exception)) def run_assertions(self): self.assertEqual(self.url.mime_type, self.mime_type) self.assertEqual(self.url.data, self.raw_data) - self.assertEqual(self.url.encoded_data, self.data) + if self.base64_encoded: + self.assertEqual(self.url.encoded_data, self.data) + else: + self.assertEqual(self.url.decoded_data, self.data) self.assertEqual(self.url.is_base64_encoded, self.base64_encoded) self.assertEqual(self.expected_url, self.url.url)