add

Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 [lib]
 name = "modern_colorthief"
 crate-type = ["cdylib"]
+path = "src/rust/lib.rs"
 
 [dependencies]
 pyo3 = "0.21.2"

docs/conf.py

@@ -5,7 +5,7 @@
 
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
-import modern_colorthief
+import python.modern_colorthief.modern_colorthief as modern_colorthief
 import datetime
 
 project = "Modern Colorthief"

examples/test.py

@@ -1,4 +1,4 @@
-from modern_colorthief import get_color, get_palette
+from python.modern_colorthief.modern_colorthief import get_color, get_palette
 from colorthief import ColorThief
 from pathlib import Path
 import os

examples/test_time.py

@@ -1,4 +1,4 @@
-from modern_colorthief import get_color, get_palette
+from python.modern_colorthief.modern_colorthief import get_color, get_palette
 from fast_colorthief import get_dominant_color, get_palette as f_get_palette
 import timeit
 from colorthief import ColorThief

pyproject.toml

@@ -27,6 +27,7 @@ fast-colorthief = "^0.0.5"
 
 [tool.maturin]
 features = ["pyo3/extension-module"]
+python-source = "src/python"
 
 [tool.poetry]
 name = "modern_colorthief"

src/python/modern_colorthief/__init__.py

@@ -0,0 +1,37 @@
+import io
+
+# Rust import
+from .modern_colorthief import *
+
+__doc__ = modern_colorthief.__doc__
+if hasattr(modern_colorthief, "__all__"):
+    __all__ = modern_colorthief.__all__
+
+
+def get_palette(
+    image: str | bytes | io.BytesIO,
+    color_count: int = 10,
+    quality: int | None = 10,
+) -> list[tuple[int, int, int]]:
+    if isinstance(image, str):
+        return _get_palette_given_string(image, color_count, quality)
+
+    if isinstance(image, bytes):
+        return _get_palette_given_bytes(image, color_count, quality)
+
+    if isinstance(image, io.BytesIO):
+        return _get_palette_given_bytes(image.getvalue(), color_count, quality)
+
+
+def get_color(
+    image: str | bytes | io.BytesIO,
+    quality: int | None = 10,
+) -> tuple[int, int, int]:
+    if isinstance(image, str):
+        return _get_color_given_location(image, quality)
+
+    if isinstance(image, bytes):
+        return _get_color_given_bytes(image, quality)
+
+    if isinstance(image, io.BytesIO):
+        return _get_color_given_bytes(image.getvalue(), quality)

src/rust/lib.rs

@@ -1,4 +1,5 @@
 use color_thief::ColorFormat;
+use image::DynamicImage;
 use pyo3::prelude::*;
 
 fn get_image_buffer(img: image::DynamicImage) -> (Vec<u8>, ColorFormat) {
@@ -10,16 +11,52 @@ fn get_image_buffer(img: image::DynamicImage) -> (Vec<u8>, ColorFormat) {
         _ => unreachable!(),
     }
 }
+/// Returns the pallette given an bytes object
 
-/// Returns the pallette given an image
 #[pyfunction]
-fn get_palette(
+fn _get_palette_given_bytes(
+    image: Vec<u8>,
+    color_count: Option<u8>,
+    quality: Option<u8>,
+) -> PyResult<Vec<(u8, u8, u8)>> {
+    let mut _image = image;
+    let img = image::load_from_memory(&_image).unwrap();
+
+    Ok(get_palette(img, color_count, quality).unwrap())
+}
+
+/// Returns the pallette given an image path
+#[pyfunction]
+fn _get_palette_given_location(
     image: String,
     color_count: Option<u8>,
     quality: Option<u8>,
 ) -> PyResult<Vec<(u8, u8, u8)>> {
     let img = image::open(&std::path::Path::new(&image)).unwrap();
-    let (buffer, color_type) = get_image_buffer(img);
+
+    Ok(get_palette(img, color_count, quality).unwrap())
+}
+
+// Gets the dominant color given an image
+#[pyfunction]
+fn _get_color_given_location(image: String, quality: Option<u8>) -> PyResult<(u8, u8, u8)> {
+    let palette = _get_palette_given_location(image, Some(5), Some(quality.unwrap_or(10))).unwrap();
+    Ok(palette[0])
+}
+
+#[pyfunction]
+fn _get_color_given_bytes(image: Vec<u8>, quality: Option<u8>) -> PyResult<(u8, u8, u8)> {
+    let palette =
+        _get_palette_given_bytes(image.to_owned(), Some(5), Some(quality.unwrap_or(10))).unwrap();
+    Ok(palette[0])
+}
+
+fn get_palette(
+    image: DynamicImage,
+    color_count: Option<u8>,
+    quality: Option<u8>,
+) -> Result<Vec<(u8, u8, u8)>, String> {
+    let (buffer, color_type) = get_image_buffer(image);
 
     let colors = color_thief::get_palette(
         &buffer,
@@ -37,13 +74,6 @@ fn get_palette(
     Ok(color_vec)
 }
 
-// Gets the dominant color given an image
-#[pyfunction]
-fn get_color(image: String, quality: Option<u8>) -> PyResult<(u8, u8, u8)> {
-    let palette = get_palette(image, Some(5), Some(quality.unwrap_or(10))).unwrap();
-    Ok(palette[0])
-}
-
 fn get_version() -> &'static str {
     static VERSION: std::sync::OnceLock<String> = std::sync::OnceLock::new();
 
@@ -53,7 +83,14 @@ fn get_version() -> &'static str {
 #[pymodule]
 fn modern_colorthief(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add("__version__", get_version())?;
-    m.add_function(wrap_pyfunction!(get_palette, m)?)?;
-    m.add_function(wrap_pyfunction!(get_color, m)?)?;
+
+    // Palette utils
+    m.add_function(wrap_pyfunction!(_get_palette_given_location, m)?)?;
+    m.add_function(wrap_pyfunction!(_get_palette_given_bytes, m)?)?;
+
+    // Color utils
+    m.add_function(wrap_pyfunction!(_get_color_given_bytes, m)?)?;
+    m.add_function(wrap_pyfunction!(_get_color_given_location, m)?)?;
+
     Ok(())
 }

test.py

@@ -1,10 +1,16 @@
 import io
 from PIL import Image
 
-img = Image.open('test.jpg', mode='r')
+img = Image.open("test.jpg", mode="r")
 
 img_byte_arr = io.BytesIO()
-img.save(img_byte_arr, format='PNG')
-img_byte_arr = img_byte_arr.getvalue()
+img.save(img_byte_arr, format="PNG")
+# img_byte_arr = img_byte_arr.getvalue()
 
-print(img_byte_arr)
+
+import modern_colorthief
+
+
+x = modern_colorthief.get_color(img_byte_arr)
+
+print(isinstance(img_byte_arr, io.BytesIO))