codearchiver/codearchiver/storage.py

import abc
import codearchiver.core
import contextlib
import logging
import os.path
import shutil
import typing


_logger = logging.getLogger(__name__)


class Storage(abc.ABC):
	@abc.abstractmethod
	def put(self, filename: str):
		'''Put a local file into storage. If an error occurs, a partial copy may remain in storage. If it completes, the local input file is removed.'''

	def put_result(self, result: codearchiver.core.Result):
		'''Put a module's Result into storage. The semantics are as for `put`, and the exact behaviour regarding partial copies and leftover files on errors is undefined.'''
		for fn in result.files:
			self.put(fn)
		for _, subresult in result.submoduleResults:
			self.put_result(subresult)

	@abc.abstractmethod
	@contextlib.contextmanager
	def open(self, filename: str) -> typing.Iterator[typing.BinaryIO]:
		'''Open a file from storage.'''


class DirectoryStorage(Storage):
	def __init__(self, directory):
		super().__init__()
		self._directory = directory

	def _check_directory(self):
		exists = os.path.exists(self._directory)
		if exists and not os.path.isdir(self._directory):
			raise NotADirectoryError(self._directory)
		return exists

	def _ensure_directory(self):
		if not self._check_directory():
			os.makedirs(self._directory)

	def put(self, filename):
		self._ensure_directory()
		_logger.info(f'Moving {filename} to {self._directory}')
		shutil.move(filename, self._directory)

	@contextlib.contextmanager
	def open(self, filename):
		with open(filename, 'rb') as fp:
			yield fp