- use tokio instead of os threads

- buffer writes
- also write in parallel

src/error.rs

@@ -6,8 +6,14 @@ pub enum Error {
     IO(#[from] std::io::Error),
     #[error(transparent)]
     ColorCet(#[from] colorcet::ColorcetError),
+    #[error(transparent)]
+    Tokio(#[from] tokio::task::JoinError),
     #[error("could not parse color: {0}")]
     ColorParse(String),
     #[error("could not covert ColorMap into LinearGradient")]
     Conversion,
+    #[error("mutex was poisoned, this is a bug, please report it!")]
+    MutexPoisoned,
+    #[error("cannot express {0} as Rational32")]
+    Rational(f64),
 }

src/lib.rs

@@ -8,20 +8,19 @@ use colorcet::ColorMap;
 use colorgrad::{Gradient, LinearGradient};
 use css_color::Srgb;
 use flate2::write::ZlibEncoder;
+use lazy_static::lazy_static;
 use ndarray::{ArcArray2, AsArray, Ix2, s};
 use num::{Complex, FromPrimitive, Rational32, traits::ToBytes};
 use rayon::prelude::*;
 use std::collections::HashSet;
 use std::fs::{File, OpenOptions};
 use std::hash::{DefaultHasher, Hash, Hasher};
-use std::io::{Read, Seek, SeekFrom, Write};
+use std::io::{BufWriter, Read, Seek, SeekFrom, Write};
 use std::path::Path;
-use std::time::Duration;
+use std::sync::{Arc, Mutex};
+use std::thread::available_parallelism;
 use std::{cmp::Ordering, collections::HashMap};
-use std::{
-    thread,
-    thread::{JoinHandle, available_parallelism, sleep},
-};
+use tokio::{runtime::Runtime, task::JoinHandle};
 use zstd::zstd_safe::CompressionLevel;
 use zstd::{DEFAULT_COMPRESSION_LEVEL, stream::Encoder};
 
@@ -29,8 +28,12 @@ const TAG_SIZE: usize = 20;
 const OFFSET_SIZE: usize = 8;
 const OFFSET: u64 = 16;
 
+lazy_static! {
+    static ref RT: Runtime = Runtime::new().unwrap();
+}
+
 /// Compression: deflate or zstd
-#[derive(Clone, Debug)]
+#[derive(Copy, Clone, Debug)]
 pub enum Compression {
     Deflate,
     Zstd(CompressionLevel),
@@ -62,32 +65,29 @@ impl IFD {
 
     fn write(
         &mut self,
-        ijtifffile: &mut IJTiffFile,
+        hashes: &mut HashMap<u64, u64>,
+        file: &mut BufWriter<File>,
         where_to_write_offset: u64,
     ) -> Result<u64, Error> {
         let mut tags = self.tags.drain().collect::<Vec<_>>();
         tags.sort();
-        ijtifffile.file.seek(SeekFrom::End(0))?;
-        if ijtifffile.file.stream_position()? % 2 == 1 {
-            ijtifffile.file.write_all(&[0])?;
+        file.seek(SeekFrom::End(0))?;
+        if file.stream_position()? % 2 == 1 {
+            file.write_all(&[0])?;
         }
-        let offset = ijtifffile.file.stream_position()?;
-        ijtifffile
-            .file
-            .write_all(&(tags.len() as u64).to_le_bytes())?;
+        let offset = file.stream_position()?;
+        file.write_all(&(tags.len() as u64).to_le_bytes())?;
 
         for tag in tags.iter_mut() {
-            tag.write_tag(ijtifffile)?;
+            tag.write_tag(file)?;
         }
-        let where_to_write_next_ifd_offset = ijtifffile.file.stream_position()?;
-        ijtifffile.file.write_all(&[0; OFFSET_SIZE])?;
+        let where_to_write_next_ifd_offset = file.stream_position()?;
+        file.write_all(&[0; OFFSET_SIZE])?;
         for tag in tags.iter() {
-            tag.write_data(ijtifffile)?;
+            tag.write_data(hashes, file)?;
         }
-        ijtifffile
-            .file
-            .seek(SeekFrom::Start(where_to_write_offset))?;
-        ijtifffile.file.write_all(&offset.to_le_bytes())?;
+        file.seek(SeekFrom::Start(where_to_write_offset))?;
+        file.write_all(&offset.to_le_bytes())?;
         Ok(where_to_write_next_ifd_offset)
     }
 }
@@ -291,13 +291,16 @@ impl Tag {
     }
 
     pub fn short_long_or_long8(code: u16, value: &[u64]) -> Self {
-        let m = *value.iter().max().unwrap();
-        if m < 65536 {
-            Tag::short(code, &value.iter().map(|x| *x as u16).collect::<Vec<_>>())
-        } else if m < 4294967296 {
-            Tag::long(code, &value.iter().map(|x| *x as u32).collect::<Vec<_>>())
+        if let Some(&m) = value.iter().max() {
+            if m < 65536 {
+                Tag::short(code, &value.iter().map(|x| *x as u16).collect::<Vec<_>>())
+            } else if m < 4294967296 {
+                Tag::long(code, &value.iter().map(|x| *x as u32).collect::<Vec<_>>())
+            } else {
+                Tag::long8(code, value)
+            }
         } else {
-            Tag::long8(code, value)
+            Tag::short(code, &[])
         }
     }
 
@@ -327,32 +330,34 @@ impl Tag {
         c as u64
     }
 
-    fn write_tag(&mut self, ijtifffile: &mut IJTiffFile) -> Result<(), Error> {
-        self.offset = ijtifffile.file.stream_position()?;
-        ijtifffile.file.write_all(&self.code.to_le_bytes())?;
-        ijtifffile.file.write_all(&self.ttype.to_le_bytes())?;
-        ijtifffile.file.write_all(&self.count().to_le_bytes())?;
+    fn write_tag(&mut self, file: &mut BufWriter<File>) -> Result<(), Error> {
+        self.offset = file.stream_position()?;
+        file.write_all(&self.code.to_le_bytes())?;
+        file.write_all(&self.ttype.to_le_bytes())?;
+        file.write_all(&self.count().to_le_bytes())?;
         if self.bytes.len() <= OFFSET_SIZE {
-            ijtifffile.file.write_all(&self.bytes)?;
-            ijtifffile
-                .file
-                .write_all(&vec![0; OFFSET_SIZE - self.bytes.len()])?;
+            file.write_all(&self.bytes)?;
+            file.write_all(&vec![0; OFFSET_SIZE - self.bytes.len()])?;
         } else {
-            ijtifffile.file.write_all(&[0; OFFSET_SIZE])?;
+            file.write_all(&[0; OFFSET_SIZE])?;
         }
         Ok(())
     }
 
-    fn write_data(&self, ijtifffile: &mut IJTiffFile) -> Result<(), Error> {
+    fn write_data(
+        &self,
+        hashes: &mut HashMap<u64, u64>,
+        file: &mut BufWriter<File>,
+    ) -> Result<(), Error> {
         if self.bytes.len() > OFFSET_SIZE {
-            ijtifffile.file.seek(SeekFrom::End(0))?;
-            let offset = ijtifffile.write(&self.bytes)?;
-            ijtifffile.file.seek(SeekFrom::Start(
+            file.seek(SeekFrom::End(0))?;
+            let offset = IJTiffFile::write(hashes, file, &self.bytes)?;
+            file.seek(SeekFrom::Start(
                 self.offset + (TAG_SIZE - OFFSET_SIZE) as u64,
             ))?;
-            ijtifffile.file.write_all(&offset.to_le_bytes())?;
-            if ijtifffile.file.stream_position()? % 2 == 1 {
-                ijtifffile.file.write_all(&[0])?;
+            file.write_all(&offset.to_le_bytes())?;
+            if file.stream_position()? % 2 == 1 {
+                file.write_all(&[0])?;
             }
         }
         Ok(())
@@ -360,18 +365,24 @@ impl Tag {
 }
 
 #[derive(Debug)]
-struct CompressedFrame {
-    bytes: Vec<Vec<u8>>,
+struct Frame {
+    offsets: Vec<u64>,
+    bytecounts: Vec<u64>,
     image_width: u32,
     image_length: u32,
-    tile_width: usize,
-    tile_length: usize,
+    tile_width: u16,
+    tile_length: u16,
     bits_per_sample: u16,
     sample_format: u16,
 }
 
-impl CompressedFrame {
-    fn new<T>(frame: ArcArray2<T>, compression: Compression) -> CompressedFrame
+impl Frame {
+    fn new<T>(
+        hashes: Arc<Mutex<HashMap<u64, u64>>>,
+        file: Arc<Mutex<BufWriter<File>>>,
+        frame: ArcArray2<T>,
+        compression: Compression,
+    ) -> Result<Frame, Error>
     where
         T: Bytes + Send + Sync,
     {
@@ -394,7 +405,7 @@ impl CompressedFrame {
                     (j + 1) * tile_length,
                 ));
             }
-            if shape[1] % tile_length != 0 {
+            if !shape[1].is_multiple_of(tile_length) {
                 slices.push((
                     i * tile_width,
                     (i + 1) * tile_width,
@@ -403,7 +414,7 @@ impl CompressedFrame {
                 ));
             }
         }
-        if shape[0] % tile_width != 0 {
+        if !shape[0].is_multiple_of(tile_width) {
             for j in 0..m {
                 slices.push((
                     n * tile_width,
@@ -412,7 +423,7 @@ impl CompressedFrame {
                     (j + 1) * tile_length,
                 ));
             }
-            if shape[1] % tile_length != 0 {
+            if !shape[1].is_multiple_of(tile_length) {
                 slices.push((n * tile_width, shape[0], m * tile_length, shape[1]));
             }
         }
@@ -423,28 +434,16 @@ impl CompressedFrame {
                     slices
                         .into_par_iter()
                         .map(|slice| {
-                            CompressedFrame::compress_tile_deflate(
-                                frame.clone(),
-                                slice,
-                                tile_size,
-                                tile_size,
-                            )
-                            .unwrap()
+                            Frame::compress_tile_deflate(frame.clone(), slice, tile_size, tile_size)
                         })
-                        .collect()
+                        .collect::<Result<Vec<_>, Error>>()?
                 } else {
                     slices
                         .into_iter()
                         .map(|slice| {
-                            CompressedFrame::compress_tile_deflate(
-                                frame.clone(),
-                                slice,
-                                tile_size,
-                                tile_size,
-                            )
-                            .unwrap()
+                            Frame::compress_tile_deflate(frame.clone(), slice, tile_size, tile_size)
                         })
-                        .collect()
+                        .collect::<Result<Vec<_>, Error>>()?
                 }
             }
 
@@ -453,43 +452,51 @@ impl CompressedFrame {
                     slices
                         .into_par_iter()
                         .map(|slice| {
-                            CompressedFrame::compress_tile_zstd(
+                            Frame::compress_tile_zstd(
                                 frame.clone(),
                                 slice,
                                 tile_size,
                                 tile_size,
                                 level,
                             )
-                            .unwrap()
                         })
-                        .collect()
+                        .collect::<Result<Vec<_>, Error>>()?
                 } else {
                     slices
                         .into_iter()
                         .map(|slice| {
-                            CompressedFrame::compress_tile_zstd(
+                            Frame::compress_tile_zstd(
                                 frame.clone(),
                                 slice,
                                 tile_size,
                                 tile_size,
                                 level,
                             )
-                            .unwrap()
                         })
-                        .collect()
+                        .collect::<Result<Vec<_>, Error>>()?
                 }
             }
         };
 
-        CompressedFrame {
-            bytes,
+        let mut offsets = Vec::new();
+        let mut bytecounts = Vec::new();
+        let mut file = file.lock().map_err(|_| Error::MutexPoisoned)?;
+        let mut hashes = hashes.lock().map_err(|_| Error::MutexPoisoned)?;
+        for tile in bytes {
+            bytecounts.push(tile.len() as u64);
+            offsets.push(IJTiffFile::write(&mut hashes, &mut file, &tile)?);
+        }
+
+        Ok(Frame {
+            offsets,
+            bytecounts,
             image_width: shape[1] as u32,
             image_length: shape[0] as u32,
-            tile_width,
-            tile_length,
+            tile_width: tile_width as u16,
+            tile_length: tile_length as u16,
             bits_per_sample: T::BITS_PER_SAMPLE,
             sample_format: T::SAMPLE_FORMAT,
-        }
+        })
     }
 
     fn encode<W, T>(
@@ -534,7 +541,7 @@ impl CompressedFrame {
         T: Bytes,
     {
         let mut encoder = ZlibEncoder::new(Vec::new(), flate2::Compression::default());
-        encoder = CompressedFrame::encode(encoder, frame, slice, tile_width, tile_length)?;
+        encoder = Frame::encode(encoder, frame, slice, tile_width, tile_length)?;
         Ok(encoder.finish()?)
     }
 
@@ -553,50 +560,13 @@ impl CompressedFrame {
         let bytes_per_sample = (T::BITS_PER_SAMPLE / 8) as usize;
         encoder.include_contentsize(true)?;
         encoder.set_pledged_src_size(Some((bytes_per_sample * tile_width * tile_length) as u64))?;
-        encoder.include_checksum(true)?;
-        encoder = CompressedFrame::encode(encoder, frame, slice, tile_width, tile_length)?;
+        encoder.include_checksum(false)?;
+        encoder = Frame::encode(encoder, frame, slice, tile_width, tile_length)?;
         encoder.finish()?;
         Ok(dest)
     }
 }
 
-#[derive(Clone, Debug)]
-struct Frame {
-    offsets: Vec<u64>,
-    bytecounts: Vec<u64>,
-    image_width: u32,
-    image_length: u32,
-    bits_per_sample: u16,
-    sample_format: u16,
-    tile_width: u16,
-    tile_length: u16,
-}
-
-impl Frame {
-    #[allow(clippy::too_many_arguments)]
-    fn new(
-        offsets: Vec<u64>,
-        bytecounts: Vec<u64>,
-        image_width: u32,
-        image_length: u32,
-        bits_per_sample: u16,
-        sample_format: u16,
-        tile_width: u16,
-        tile_length: u16,
-    ) -> Self {
-        Frame {
-            offsets,
-            bytecounts,
-            image_width,
-            image_length,
-            bits_per_sample,
-            sample_format,
-            tile_width,
-            tile_length,
-        }
-    }
-}
-
 /// trait to convert numbers to bytes
 pub trait Bytes {
     const BITS_PER_SAMPLE: u16;
@@ -654,10 +624,10 @@ pub enum Colors {
 /// save 2d arrays in a tif file compatible with Fiji/ImageJ
 #[derive(Debug)]
 pub struct IJTiffFile {
-    file: File,
+    file: Arc<Mutex<BufWriter<File>>>,
     frames: HashMap<(usize, usize, usize), Frame>,
-    hashes: HashMap<u64, u64>,
-    threads: HashMap<(usize, usize, usize), JoinHandle<CompressedFrame>>,
+    hashes: Arc<Mutex<HashMap<u64, u64>>>,
+    threads: HashMap<(usize, usize, usize), JoinHandle<Result<Frame, Error>>>,
     /// zstd: -7 ..= 22
     pub compression: Compression,
     pub colors: Colors,
@@ -681,7 +651,7 @@ impl Drop for IJTiffFile {
 }
 
 impl IJTiffFile {
-    /// create new tifffile from path, use it's save() method to save frames
+    /// create new tifffile from path, use its save() method to save frames
     /// the file is finalized when it goes out of scope
     pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
         let mut file = OpenOptions::new()
@@ -696,9 +666,9 @@ impl IJTiffFile {
         file.write_all(&0u16.to_le_bytes())?;
         file.write_all(&OFFSET.to_le_bytes())?;
         Ok(IJTiffFile {
-            file,
+            file: Arc::new(Mutex::new(BufWriter::new(file))),
             frames: HashMap::new(),
-            hashes: HashMap::new(),
+            hashes: Arc::new(Mutex::new(HashMap::new())),
             threads: HashMap::new(),
             compression: Compression::Zstd(DEFAULT_COMPRESSION_LEVEL),
             colors: Colors::None,
@@ -842,29 +812,31 @@ impl IJTiffFile {
         hasher.finish()
     }
 
-    fn hash_check(&mut self, bytes: &Vec<u8>, offset: u64) -> Result<bool, Error> {
-        let current_offset = self.file.stream_position()?;
-        self.file.seek(SeekFrom::Start(offset))?;
+    fn hash_check(f: &mut BufWriter<File>, bytes: &Vec<u8>, offset: u64) -> Result<bool, Error> {
+        let current_offset = f.stream_position()?;
+        f.seek(SeekFrom::Start(offset))?;
         let mut buffer = vec![0; bytes.len()];
-        self.file.read_exact(&mut buffer)?;
+        f.get_ref().read_exact(&mut buffer)?;
         let same = bytes == &buffer;
-        self.file.seek(SeekFrom::Start(current_offset))?;
+        f.seek(SeekFrom::Start(current_offset))?;
         Ok(same)
     }
 
-    fn write(&mut self, bytes: &Vec<u8>) -> Result<u64, Error> {
+    fn write(
+        hashes: &mut HashMap<u64, u64>,
+        file: &mut BufWriter<File>,
+        bytes: &Vec<u8>,
+    ) -> Result<u64, Error> {
         let hash = IJTiffFile::hash(&bytes);
-        if self.hashes.contains_key(&hash)
-            && self.hash_check(bytes, *self.hashes.get(&hash).unwrap())?
-        {
-            Ok(*self.hashes.get(&hash).unwrap())
+        if hashes.contains_key(&hash) && Self::hash_check(file, bytes, hashes[&hash])? {
+            Ok(hashes[&hash])
         } else {
-            if self.file.stream_position()? % 2 == 1 {
-                self.file.write_all(&[0])?;
+            if file.stream_position()? % 2 == 1 {
+                file.write_all(&[0])?;
             }
-            let offset = self.file.stream_position()?;
-            self.hashes.insert(hash, offset);
-            self.file.write_all(bytes)?;
+            let offset = file.stream_position()?;
+            hashes.insert(hash, offset);
+            file.write_all(bytes)?;
             Ok(offset)
         }
     }
@@ -875,58 +847,31 @@ impl IJTiffFile {
         A: AsArray<'a, T, Ix2>,
         T: Bytes + Clone + Send + Sync + 'static,
     {
-        let n_threads = usize::from(available_parallelism()?);
-        loop {
-            self.collect_threads(false)?;
-            if self.threads.len() < n_threads {
-                break;
-            }
-            sleep(Duration::from_millis(100));
-        }
-        let compression = self.compression.clone();
+        self.collect_threads(false, usize::from(available_parallelism()?))?;
+        let hashes = self.hashes.clone();
+        let file = self.file.clone();
         let frame = frame.into().to_shared();
+        let compression = self.compression;
         self.threads.insert(
             (c, z, t),
-            thread::spawn(move || CompressedFrame::new(frame, compression)),
+            RT.spawn_blocking(move || Frame::new(hashes, file, frame, compression)),
         );
         Ok(())
     }
 
-    fn collect_threads(&mut self, block: bool) -> Result<(), Error> {
-        for (c, z, t) in self.threads.keys().cloned().collect::<Vec<_>>() {
-            if block || self.threads[&(c, z, t)].is_finished() {
-                if let Some(thread) = self.threads.remove(&(c, z, t)) {
-                    self.write_frame(thread.join().unwrap(), c, z, t)?;
+    fn collect_threads(&mut self, block: bool, max_threads: usize) -> Result<(), Error> {
+        RT.block_on(async {
+            while self.threads.len() > max_threads {
+                for (c, z, t) in self.threads.keys().cloned().collect::<Vec<_>>() {
+                    if (block || self.threads[&(c, z, t)].is_finished())
+                        && let Some(thread) = self.threads.remove(&(c, z, t))
+                    {
+                        self.frames.insert((c, z, t), thread.await??);
+                    }
                 }
             }
-        }
-        Ok(())
-    }
-
-    fn write_frame(
-        &mut self,
-        frame: CompressedFrame,
-        c: usize,
-        z: usize,
-        t: usize,
-    ) -> Result<(), Error> {
-        let mut offsets = Vec::new();
-        let mut bytecounts = Vec::new();
-        for tile in frame.bytes {
-            bytecounts.push(tile.len() as u64);
-            offsets.push(self.write(&tile)?);
-        }
-        let frame = Frame::new(
-            offsets,
-            bytecounts,
-            frame.image_width,
-            frame.image_length,
-            frame.bits_per_sample,
-            frame.sample_format,
-            frame.tile_width as u16,
-            frame.tile_length as u16,
-        );
-        self.frames.insert((c, z, t), frame);
+            Ok::<(), Error>(())
+        })?;
         Ok(())
     }
 
@@ -957,7 +902,7 @@ impl IJTiffFile {
     }
 
     fn close(&mut self) -> Result<(), Error> {
-        self.collect_threads(true)?;
+        self.collect_threads(true, 0)?;
         let mut c_size = 1;
         let mut z_size = 1;
         let mut t_size = 1;
@@ -970,6 +915,8 @@ impl IJTiffFile {
         let mut where_to_write_next_ifd_offset = OFFSET - OFFSET_SIZE as u64;
         let mut warn = Vec::new();
         let (samples_per_pixel, n_frames) = self.spp_and_n_frames(c_size, t_size, z_size);
+        let mut file = self.file.lock().map_err(|_| Error::MutexPoisoned)?;
+        let mut hashes = self.hashes.lock().map_err(|_| Error::MutexPoisoned)?;
         for frame_number in 0..n_frames {
             if let Some(frame) = self
                 .frames
@@ -1009,7 +956,7 @@ impl IJTiffFile {
                     ifd.tags.insert(Tag::short(339, &[frame.sample_format]));
                 }
                 if let Some(px_size) = self.px_size {
-                    let r = [Rational32::from_f64(px_size).unwrap()];
+                    let r = [Rational32::from_f64(px_size).ok_or_else(|| Error::Rational(px_size))?];
                     ifd.tags.insert(Tag::rational(282, &r));
                     ifd.tags.insert(Tag::rational(283, &r));
                     ifd.tags.insert(Tag::short(296, &[1]));
@@ -1019,27 +966,27 @@ impl IJTiffFile {
                 } else if let Colors::None = self.colors {
                     ifd.tags.insert(Tag::short(262, &[1]));
                 }
-                if frame_number == 0 {
-                    if let Colors::Colormap(colormap) = &self.colors {
-                        ifd.tags.insert(Tag::short(
-                            320,
-                            &self.get_colormap(colormap, frame.bits_per_sample),
-                        ));
-                    }
+                if frame_number == 0
+                    && let Colors::Colormap(colormap) = &self.colors
+                {
+                    ifd.tags.insert(Tag::short(
+                        320,
+                        &self.get_colormap(colormap, frame.bits_per_sample),
+                    ));
                 }
-                if frame_number < c_size {
-                    if let Colors::Colors(colors) = &self.colors {
-                        ifd.tags.insert(Tag::short(
-                            320,
-                            &self.get_color(&colors[frame_number], frame.bits_per_sample),
-                        ));
-                        ifd.tags.insert(Tag::short(262, &[3]));
-                    }
+                if frame_number < c_size
+                    && let Colors::Colors(colors) = &self.colors
+                {
+                    ifd.tags.insert(Tag::short(
+                        320,
+                        &self.get_color(&colors[frame_number], frame.bits_per_sample),
+                    ));
+                    ifd.tags.insert(Tag::short(262, &[3]));
                 }
-                if let Colors::None = &self.colors {
-                    if c_size > 1 {
-                        ifd.tags.insert(Tag::short(284, &[2]));
-                    }
+                if let Colors::None = &self.colors
+                    && c_size > 1
+                {
+                    ifd.tags.insert(Tag::short(284, &[2]));
                 }
                 for channel in 0..samples_per_pixel {
                     let czt = self.get_czt(frame_number, channel, c_size, z_size);
@@ -1060,7 +1007,8 @@ impl IJTiffFile {
                         &format!("{}", Utc::now().format("%Y:%m:%d %H:%M:%S")),
                     ));
                 }
-                where_to_write_next_ifd_offset = ifd.write(self, where_to_write_next_ifd_offset)?;
+                where_to_write_next_ifd_offset =
+                    ifd.write(&mut hashes, &mut file, where_to_write_next_ifd_offset)?;
             } else {
                 warn.push((frame_number, 0));
             }
@@ -1076,9 +1024,8 @@ impl IJTiffFile {
                 )
             }
         }
-        self.file
-            .seek(SeekFrom::Start(where_to_write_next_ifd_offset))?;
-        self.file.write_all(&0u64.to_le_bytes())?;
+        file.seek(SeekFrom::Start(where_to_write_next_ifd_offset))?;
+        file.write_all(&0u64.to_le_bytes())?;
         Ok(())
     }
 }

src/py.rs

@@ -10,8 +10,7 @@ impl From<crate::error::Error> for PyErr {
     }
 }
 
-#[pyclass(subclass)]
-#[pyo3(name = "Tag")]
+#[pyclass(name = "Tag", module = "tiffwrite_rs", subclass, from_py_object)]
 #[derive(Clone, Debug)]
 struct PyTag {
     tag: Tag,
@@ -163,8 +162,7 @@ impl PyTag {
     }
 }
 
-#[pyclass(subclass)]
-#[pyo3(name = "IJTiffFile")]
+#[pyclass(name = "IJTiffFile", module = "tiffwrite_rs", subclass)]
 #[derive(Debug)]
 struct PyIJTiffFile {
     ijtifffile: Option<IJTiffFile>,
@@ -199,10 +197,10 @@ impl PyIJTiffFile {
 
     #[getter]
     fn get_colors(&self) -> PyResult<Option<Vec<Vec<u8>>>> {
-        if let Some(ijtifffile) = &self.ijtifffile {
-            if let Colors::Colors(colors) = &ijtifffile.colors {
-                return Ok(Some(colors.to_owned()));
-            }
+        if let Some(ijtifffile) = &self.ijtifffile
+            && let Colors::Colors(colors) = &ijtifffile.colors
+        {
+            return Ok(Some(colors.to_owned()));
         }
         Ok(None)
     }
@@ -217,10 +215,10 @@ impl PyIJTiffFile {
 
     #[getter]
     fn get_colormap(&mut self) -> PyResult<Option<Vec<Vec<u8>>>> {
-        if let Some(ijtifffile) = &self.ijtifffile {
-            if let Colors::Colormap(colormap) = &ijtifffile.colors {
-                return Ok(Some(colormap.to_owned()));
-            }
+        if let Some(ijtifffile) = &self.ijtifffile
+            && let Colors::Colormap(colormap) = &ijtifffile.colors
+        {
+            return Ok(Some(colormap.to_owned()));
         }
         Ok(None)
     }
@@ -303,25 +301,25 @@ impl PyIJTiffFile {
 
     #[pyo3(signature = (tag, czt=None))]
     fn append_extra_tag(&mut self, tag: PyTag, czt: Option<(usize, usize, usize)>) {
-        if let Some(ijtifffile) = self.ijtifffile.as_mut() {
-            if let Some(extra_tags) = ijtifffile.extra_tags.get_mut(&czt) {
-                extra_tags.push(tag.tag)
-            }
+        if let Some(ijtifffile) = self.ijtifffile.as_mut()
+            && let Some(extra_tags) = ijtifffile.extra_tags.get_mut(&czt)
+        {
+            extra_tags.push(tag.tag)
         }
     }
 
     #[pyo3(signature = (czt=None))]
     fn get_tags(&self, czt: Option<(usize, usize, usize)>) -> PyResult<Vec<PyTag>> {
-        if let Some(ijtifffile) = &self.ijtifffile {
-            if let Some(extra_tags) = ijtifffile.extra_tags.get(&czt) {
-                let v = extra_tags
-                    .iter()
-                    .map(|tag| PyTag {
-                        tag: tag.to_owned(),
-                    })
-                    .collect();
-                return Ok(v);
-            }
+        if let Some(ijtifffile) = &self.ijtifffile
+            && let Some(extra_tags) = ijtifffile.extra_tags.get(&czt)
+        {
+            let v = extra_tags
+                .iter()
+                .map(|tag| PyTag {
+                    tag: tag.to_owned(),
+                })
+                .collect();
+            return Ok(v);
         }
         Ok(Vec::new())
     }
@@ -369,9 +367,17 @@ impl_save! {
 
 #[pymodule]
 #[pyo3(name = "tiffwrite_rs")]
-fn tiffwrite_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    color_eyre::install()?;
-    m.add_class::<PyTag>()?;
-    m.add_class::<PyIJTiffFile>()?;
-    Ok(())
+mod tiffwrite_rs {
+    use pyo3::prelude::*;
+
+    #[pymodule_export]
+    use super::PyTag;
+
+    #[pymodule_export]
+    use super::PyIJTiffFile;
+
+    #[pymodule_init]
+    fn init(_: &Bound<'_, PyModule>) -> PyResult<()> {
+        Ok(color_eyre::install()?)
+    }
 }