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

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

2

3 from __future__ import annotations

4

5 from enum import Enum

6

7

8 class 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

53 class 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)