Coverage for src/polystore/formats.py: 73%

1""" 

2File format definitions for polystore. 

3 

4This module defines the supported file formats and their extensions. 

5""" 

6 

7from enum import Enum 

8 

9 

10class FileFormat(Enum): 

11 """Enumeration of supported file formats.""" 

12 

13 # Array formats 

14 NUMPY = "numpy" 

15 TORCH = "torch" 

16 JAX = "jax" 

17 CUPY = "cupy" 

18 TENSORFLOW = "tensorflow" 

19 ZARR = "zarr" 

20 

21 # Image formats 

22 TIFF = "tiff" 

23 

24 # Data formats 

25 CSV = "csv" 

26 JSON = "json" 

27 TEXT = "text" 

28 

29 # ROI format 

30 ROI = "roi" 

31 

32 @property 

33 def extensions(self): 

34 """Get file extensions for this format.""" 

35 return FILE_FORMAT_EXTENSIONS.get(self, []) 

36 

37 

38# Mapping of file formats to their extensions 

39FILE_FORMAT_EXTENSIONS = { 

40 FileFormat.NUMPY: [".npy", ".npz"], 

41 FileFormat.TORCH: [".pt", ".pth"], 

42 FileFormat.JAX: [".jax"], 

43 FileFormat.CUPY: [".cupy"], 

44 FileFormat.TENSORFLOW: [".tf"], 

45 FileFormat.ZARR: [".zarr"], 

46 FileFormat.TIFF: [".tif", ".tiff"], 

47 FileFormat.CSV: [".csv"], 

48 FileFormat.JSON: [".json"], 

49 FileFormat.TEXT: [".txt"], 

50 FileFormat.ROI: [".roi.zip"], 

51} 

52 

53# Default image extensions 

54DEFAULT_IMAGE_EXTENSIONS = {".tif", ".tiff", ".png", ".jpg", ".jpeg"} 

55 

56 

57def get_format_from_extension(ext: str) -> FileFormat: 

58 """ 

59 Get file format from extension. 

60 

61 Args: 

62 ext: File extension (with or without leading dot) 

63 

64 Returns: 

65 FileFormat enum value 

66 

67 Raises: 

68 ValueError: If extension is not recognized 

69 """ 

70 if not ext.startswith("."): 

71 ext = f".{ext}" 

72 

73 ext = ext.lower() 

74 

75 for fmt, extensions in FILE_FORMAT_EXTENSIONS.items(): 

76 if ext in extensions: 

77 return fmt 

78 

79 raise ValueError(f"Unknown file extension: {ext}")