diff --git a/src/loadimg/loadimg.py b/src/loadimg/loadimg.py index 3162717..57d592a 100644 --- a/src/loadimg/loadimg.py +++ b/src/loadimg/loadimg.py @@ -14,7 +14,7 @@ def main(): parser.add_argument("input", help="Input image (file path, URL, or base64 string)") parser.add_argument( "--output-type", - choices=["pil", "numpy", "str", "base64", "ascii", "ansi"], + choices=["pil", "numpy", "str", "base64", "ascii", "ansi", "url"], default="ansi", help="Output format (default: ansi)", ) diff --git a/src/loadimg/utils.py b/src/loadimg/utils.py index f5095db..12bd492 100644 --- a/src/loadimg/utils.py +++ b/src/loadimg/utils.py @@ -11,6 +11,7 @@ from concurrent.futures import ThreadPoolExecutor import glob from tqdm import tqdm +import json # TODO: # support other input types such as lists, tensors, ... @@ -18,7 +19,7 @@ def load_img( img: Union[str, bytes, np.ndarray, Image.Image], - output_type: Literal["pil", "numpy", "str", "base64", "ascii", "ansi"] = "pil", + output_type: Literal["pil", "numpy", "str", "base64", "ascii", "ansi", "url"] = "pil", input_type: Literal["auto", "base64", "file", "url", "numpy", "pil"] = "auto", ) -> Any: """Loads an image from various sources and returns it in a specified format. @@ -28,7 +29,8 @@ def load_img( a NumPy array, or a Pillow Image object. output_type: The desired output type. Can be "pil" (Pillow Image), "numpy" (NumPy array), "str" (file path), "base64" (base64 string), - "ascii" (ASCII art), or "ansi" (ANSI art). + "ascii" (ASCII art), "ansi" (ANSI art), or "url" (a public URL + after uploading to a temporary hosting service). input_type: The type of the input image. If set to "auto", the function will try to automatically determine the type. @@ -42,6 +44,11 @@ def load_img( # Convert to ANSI art ansi_art = load_img("image.png", output_type="ansi") + + # Upload an image and get a temporary URL + # Note: This requires an active internet connection. + # url = load_img("image.png", output_type="url") + # print(f"Image available at: {url}") ``` """ try: @@ -68,6 +75,26 @@ def load_img( img.save(buffer, format=img_type) img_str = base64.b64encode(buffer.getvalue()).decode("utf-8") return f"data:image/{img_type.lower()};base64,{img_str}" + elif output_type == "url": + upload_url = "https://uguu.se/upload" + with BytesIO() as buffer: + img_format = img.format or "PNG" + img.save(buffer, format=img_format) + buffer.seek(0) + + file_name = original_name or f"{uuid.uuid4()}.{img_format.lower()}" + + files = {'files[]': (file_name, buffer.getvalue(), f'image/{img_format.lower()}')} + + try: + # Added a timeout for robustness + response = requests.post(upload_url, files=files, timeout=15) + response.raise_for_status() + # The response body from uguu.se is expected to be the direct URL + reply = response.text.strip() + return json.loads(reply)["files"][0]["url"] + except requests.exceptions.RequestException as e: + raise IOError(f"Failed to upload image to {upload_url}: {e}") from e elif output_type == "ascii": return image_to_ascii(img) elif output_type == "ansi": @@ -81,7 +108,7 @@ def load_img( def load_imgs( imgs: Union[str, List[Union[str, bytes, np.ndarray, Image.Image]]], - output_type: Literal["pil", "numpy", "str", "base64", "ascii", "ansi"] = "pil", + output_type: Literal["pil", "numpy", "str", "base64", "ascii", "ansi", "url"] = "pil", input_type: Literal["auto", "base64", "file", "url", "numpy", "pil"] = "auto", max_workers: int = 1, glob_pattern: str = "*", @@ -93,7 +120,7 @@ def load_imgs( - Directory path (str) - List of image sources - Glob pattern - output_type: The desired output type + output_type: The desired output type. Can be "pil", "numpy", "str", "base64", "ascii", "ansi", or "url". input_type: The type of input images max_workers: Max number of parallel workers glob_pattern: Pattern for filtering files when imgs is a directory diff --git a/tests/test_cli.py b/tests/test_cli.py index 18880d3..9047a1f 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -7,6 +7,23 @@ class TestCLI(unittest.TestCase): + @patch("loadimg.utils.requests.post") + @patch("argparse.ArgumentParser.parse_args") + def test_main_with_url_output_type(self, mock_args, mock_post): + # Mock upload response + mock_response = MagicMock() + mock_response.text = '{"files": [{"url": "https://fake.uguu.se/fake.png"}]}' + mock_response.raise_for_status.return_value = None + mock_post.return_value = mock_response + + mock_args.return_value = MagicMock( + input=self.test_image_path, output_type="url", input_type="auto" + ) + + with patch("builtins.print") as mock_print: + exit_code = main() + self.assertEqual(exit_code, 0) + mock_print.assert_called() def setUp(self): # Create a temporary test image self.temp_dir = tempfile.TemporaryDirectory() diff --git a/tests/test_loadimg.py b/tests/test_loadimg.py index 0c6646b..06da913 100644 --- a/tests/test_loadimg.py +++ b/tests/test_loadimg.py @@ -50,6 +50,17 @@ def test_load_img_from_file(self): self.assertIn("\x1b[48;2;", ansi_art) self.assertIn("\x1b[0m", ansi_art) + @patch("loadimg.utils.requests.post") + def test_load_img_from_file_url(self, mock_post): + # Mock the response from the upload service + mock_response = MagicMock() + mock_response.text = '{"files": [{"url": "https://fake.uguu.se/fake.png"}]}' + mock_response.raise_for_status.return_value = None + mock_post.return_value = mock_response + + url = load_img(self.sample_image_path, output_type="url") + self.assertEqual(url, "https://fake.uguu.se/fake.png") + def test_load_img_from_base64(self): # Existing tests img = load_img(f"data:image/png;base64,{self.sample_base64}", output_type="pil") @@ -81,6 +92,16 @@ def test_load_img_from_base64(self): self.assertIn("\x1b[48;2;", ansi_art) self.assertIn("\x1b[0m", ansi_art) + @patch("loadimg.utils.requests.post") + def test_load_img_from_base64_url(self, mock_post): + mock_response = MagicMock() + mock_response.text = '{"files": [{"url": "https://fake.uguu.se/fake.png"}]}' + mock_response.raise_for_status.return_value = None + mock_post.return_value = mock_response + + url = load_img(f"data:image/png;base64,{self.sample_base64}", output_type="url") + self.assertEqual(url, "https://fake.uguu.se/fake.png") + def test_load_img_from_numpy(self): # Existing tests img = load_img(self.sample_numpy_array, output_type="pil") @@ -102,6 +123,16 @@ def test_load_img_from_numpy(self): ansi_art = load_img(self.sample_numpy_array, output_type="ansi") self.assertIn("48;2;", ansi_art) + @patch("loadimg.utils.requests.post") + def test_load_img_from_numpy_url(self, mock_post): + mock_response = MagicMock() + mock_response.text = '{"files": [{"url": "https://fake.uguu.se/fake.png"}]}' + mock_response.raise_for_status.return_value = None + mock_post.return_value = mock_response + + url = load_img(self.sample_numpy_array, output_type="url") + self.assertEqual(url, "https://fake.uguu.se/fake.png") + def test_load_img_from_pil(self): # Existing tests img = load_img(self.sample_image, output_type="pil") @@ -123,6 +154,16 @@ def test_load_img_from_pil(self): ansi_art = load_img(self.sample_image, output_type="ansi") self.assertTrue(ansi_art.count("\x1b[0m") >= 50) + @patch("loadimg.utils.requests.post") + def test_load_img_from_pil_url(self, mock_post): + mock_response = MagicMock() + mock_response.text = '{"files": [{"url": "https://fake.uguu.se/fake.png"}]}' + mock_response.raise_for_status.return_value = None + mock_post.return_value = mock_response + + url = load_img(self.sample_image, output_type="url") + self.assertEqual(url, "https://fake.uguu.se/fake.png") + @patch("requests.get") def test_load_img_from_url(self, mock_get): # Mock setup @@ -152,6 +193,24 @@ def test_load_img_from_url(self, mock_get): ansi_art = load_img("https://example.com/sample.png", output_type="ansi") self.assertIn("\x1b[48;2;", ansi_art) + @patch("loadimg.utils.requests.post") + @patch("requests.get") + def test_load_img_from_url_url(self, mock_get, mock_post): + # Mock download + mock_response_get = MagicMock() + mock_response_get.content = open(self.sample_image_path, "rb").read() + mock_response_get.raise_for_status.return_value = None + mock_get.return_value = mock_response_get + + # Mock upload + mock_response_post = MagicMock() + mock_response_post.text = '{"files": [{"url": "https://fake.uguu.se/fake.png"}]}' + mock_response_post.raise_for_status.return_value = None + mock_post.return_value = mock_response_post + + url = load_img("https://example.com/sample.png", output_type="url") + self.assertEqual(url, "https://fake.uguu.se/fake.png") + # New tests for validate_image def test_validate_image_non_pil(self): # Test with a non-PIL input