StreamWrapper is introduced to wrap file handler generated by DataPipe operation like `FileOpener`. StreamWrapper would guarantee the wrapped file handler is closed when it's out of scope.
| 288 | |
| 289 | |
| 290 | class StreamWrapper: |
| 291 | """ |
| 292 | StreamWrapper is introduced to wrap file handler generated by DataPipe operation like `FileOpener`. |
| 293 | |
| 294 | StreamWrapper would guarantee the wrapped file handler is closed when it's out of scope. |
| 295 | """ |
| 296 | |
| 297 | session_streams: Dict[Any, int] = {} |
| 298 | debug_unclosed_streams: bool = False |
| 299 | |
| 300 | def __init__(self, file_obj, parent_stream=None, name=None): |
| 301 | self.file_obj = file_obj |
| 302 | self.child_counter = 0 |
| 303 | self.parent_stream = parent_stream |
| 304 | self.close_on_last_child = False |
| 305 | self.name = name |
| 306 | self.closed = False |
| 307 | if parent_stream is not None: |
| 308 | if not isinstance(parent_stream, StreamWrapper): |
| 309 | raise RuntimeError(f'Parent stream should be StreamWrapper, {type(parent_stream)} was given') |
| 310 | parent_stream.child_counter += 1 |
| 311 | self.parent_stream = parent_stream |
| 312 | if StreamWrapper.debug_unclosed_streams: |
| 313 | StreamWrapper.session_streams[self] = 1 |
| 314 | |
| 315 | @classmethod |
| 316 | def close_streams(cls, v, depth=0): |
| 317 | """Traverse structure and attempts to close all found StreamWrappers on best effort basis.""" |
| 318 | if depth > 10: |
| 319 | return |
| 320 | if isinstance(v, StreamWrapper): |
| 321 | v.close() |
| 322 | else: |
| 323 | # Traverse only simple structures |
| 324 | if isinstance(v, dict): |
| 325 | for vv in v.values(): |
| 326 | cls.close_streams(vv, depth=depth + 1) |
| 327 | elif isinstance(v, (list, tuple)): |
| 328 | for vv in v: |
| 329 | cls.close_streams(vv, depth=depth + 1) |
| 330 | |
| 331 | def __getattr__(self, name): |
| 332 | file_obj = self.__dict__['file_obj'] |
| 333 | return getattr(file_obj, name) |
| 334 | |
| 335 | def close(self, *args, **kwargs): |
| 336 | if self.closed: |
| 337 | return |
| 338 | if StreamWrapper.debug_unclosed_streams: |
| 339 | del StreamWrapper.session_streams[self] |
| 340 | if hasattr(self, "parent_stream") and self.parent_stream is not None: |
| 341 | self.parent_stream.child_counter -= 1 |
| 342 | if not self.parent_stream.child_counter and self.parent_stream.close_on_last_child: |
| 343 | self.parent_stream.close() |
| 344 | try: |
| 345 | self.file_obj.close(*args, **kwargs) |
| 346 | except AttributeError: |
| 347 | pass |
no outgoing calls
searching dependent graphs…