Coverage for src / crump / file_types.py: 97%

1"""File type enumerations for Crump.""" 

2 

3from __future__ import annotations 

4 

5from enum import Enum 

6 

7 

8class OutputFileType(Enum): 

9 """Supported output file formats for data extraction.""" 

10 

11 CSV = "csv" 

12 PARQUET = "parquet" 

13 

14 @classmethod 

15 def from_extension(cls, extension: str) -> OutputFileType: 

16 """Determine output file type from file extension. 

17 

18 Args: 

19 extension: File extension (with or without leading dot) 

20 

21 Returns: 

22 OutputFileType enum value 

23 

24 Raises: 

25 ValueError: If extension is not supported 

26 """ 

27 ext = extension.lower().lstrip(".") 

28 if ext == "csv": 

29 return cls.CSV 

30 elif ext in ["parquet", "pq"]: 

31 return cls.PARQUET 

32 else: 

33 raise ValueError(f"Unsupported output file extension: {extension}") 

34 

35 @classmethod 

36 def from_path(cls, path: str) -> OutputFileType: 

37 """Determine output file type from file path. 

38 

39 Args: 

40 path: File path 

41 

42 Returns: 

43 OutputFileType enum value 

44 

45 Raises: 

46 ValueError: If path extension is not supported 

47 """ 

48 from pathlib import Path 

49 

50 return cls.from_extension(Path(path).suffix) 

51 

52 

53class InputFileType(Enum): 

54 """Supported input file formats for data processing.""" 

55 

56 CSV = "csv" 

57 PARQUET = "parquet" 

58 CDF = "cdf" 

59 

60 @classmethod 

61 def from_extension(cls, extension: str, default: InputFileType | None = None) -> InputFileType: 

62 """Determine input file type from file extension. 

63 

64 Args: 

65 extension: File extension (with or without leading dot) 

66 default: Default value to return if extension not recognized 

67 

68 Returns: 

69 InputFileType enum value 

70 

71 Raises: 

72 ValueError: If extension is not supported and no default provided 

73 """ 

74 ext = extension.lower().lstrip(".") 

75 if ext == "csv": 

76 return cls.CSV 

77 elif ext in ["parquet", "pq"]: 

78 return cls.PARQUET 

79 elif ext == "cdf": 

80 return cls.CDF 

81 elif default is not None: 

82 return default 

83 else: 

84 raise ValueError(f"Unsupported input file extension: {extension}") 

85 

86 @classmethod 

87 def from_path(cls, path: str, default: InputFileType | None = None) -> InputFileType: 

88 """Determine input file type from file path. 

89 

90 Args: 

91 path: File path 

92 default: Default value to return if extension not recognized 

93 

94 Returns: 

95 InputFileType enum value 

96 

97 Raises: 

98 ValueError: If path extension is not supported and no default provided 

99 """ 

100 from pathlib import Path 

101 

102 return cls.from_extension(Path(path).suffix, default=default)