Coverage for src / cufile_patcher / cufile.py: 100%

1from __future__ import annotations 

2 

3import ctypes 

4import os 

5 

6from .bindings import ( 

7 cuFileDriverClose, 

8 cuFileDriverOpen, 

9 cuFileHandleDeregister, 

10 cuFileHandleRegister, 

11 cuFileRead, 

12 cuFileWrite, 

13) 

14 

15 

16class CuFileDriver: 

17 _instance: CuFileDriver | None = None 

18 

19 def __new__(cls) -> CuFileDriver: 

20 if cls._instance is None: 

21 cls._instance = super().__new__(cls) 

22 cls._instance._initialized = False 

23 return cls._instance 

24 

25 def __init__(self) -> None: 

26 if self._initialized: 

27 return 

28 cuFileDriverOpen() 

29 self._initialized = True 

30 

31 def close(self) -> None: 

32 if self._initialized: 

33 cuFileDriverClose() 

34 self._initialized = False 

35 

36 def __del__(self) -> None: 

37 self.close() 

38 

39 @classmethod 

40 def _reset_for_tests(cls) -> None: 

41 if cls._instance is not None: 

42 cls._instance.close() 

43 cls._instance = None 

44 

45 

46def _os_mode(mode: str) -> int: 

47 modes = { 

48 "r": os.O_RDONLY, 

49 "r+": os.O_RDWR, 

50 "w": os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 

51 "w+": os.O_CREAT | os.O_RDWR | os.O_TRUNC, 

52 "a": os.O_CREAT | os.O_WRONLY | os.O_APPEND, 

53 "a+": os.O_CREAT | os.O_RDWR | os.O_APPEND, 

54 } 

55 try: 

56 return modes[mode] 

57 except KeyError as exc: 

58 raise ValueError(f"unsupported mode: {mode}") from exc 

59 

60 

61class CuFile: 

62 def __init__(self, path: str, mode: str = "r", use_direct_io: bool = False): 

63 self._driver = CuFileDriver() 

64 self._path = path 

65 self._mode = mode 

66 self._os_mode = _os_mode(mode) 

67 if use_direct_io and hasattr(os, "O_DIRECT"): 

68 self._os_mode |= os.O_DIRECT 

69 

70 self._handle: int | None = None 

71 self._cu_file_handle = None 

72 

73 @property 

74 def is_open(self) -> bool: 

75 return self._handle is not None 

76 

77 def open(self) -> None: 

78 if self.is_open: 

79 return 

80 self._handle = os.open(self._path, self._os_mode) 

81 self._cu_file_handle = cuFileHandleRegister(self._handle) 

82 

83 def close(self) -> None: 

84 if not self.is_open: 

85 return 

86 cuFileHandleDeregister(self._cu_file_handle) 

87 os.close(self._handle) 

88 self._handle = None 

89 self._cu_file_handle = None 

90 

91 def __enter__(self) -> CuFile: 

92 self.open() 

93 return self 

94 

95 def __exit__(self, exc_type, exc_val, exc_tb) -> None: 

96 self.close() 

97 

98 def __del__(self) -> None: 

99 self.close() 

100 

101 def read( 

102 self, 

103 dest: ctypes.c_void_p, 

104 size: int, 

105 file_offset: int = 0, 

106 dev_offset: int = 0, 

107 ) -> int: 

108 if not self.is_open: 

109 raise OSError("File is not open.") 

110 return cuFileRead(self._cu_file_handle, dest, size, file_offset, dev_offset) 

111 

112 def write( 

113 self, 

114 src: ctypes.c_void_p, 

115 size: int, 

116 file_offset: int = 0, 

117 dev_offset: int = 0, 

118 ) -> int: 

119 if not self.is_open: 

120 raise OSError("File is not open.") 

121 return cuFileWrite(self._cu_file_handle, src, size, file_offset, dev_offset) 

122 

123 def get_handle(self) -> int | None: 

124 return self._handle