Use pathclass's assert_disk_space.

This commit is contained in:
voussoir 2022-01-19 21:01:43 -08:00
parent ddfead3ecb
commit 822f01b56e
No known key found for this signature in database
GPG key ID: 5F7554F8C26DACCB

View file

@ -13,54 +13,25 @@ drive:
Filepath to the drive you want to check. Defaults to cwd drive. Filepath to the drive you want to check. Defaults to cwd drive.
''' '''
import argparse import argparse
import os
import shutil
import sys import sys
from voussoirkit import betterhelp from voussoirkit import betterhelp
from voussoirkit import bytestring from voussoirkit import bytestring
from voussoirkit import dotdict from voussoirkit import pathclass
from voussoirkit import vlogging from voussoirkit import vlogging
log = vlogging.getLogger(__name__, 'reserve_disk_space') log = vlogging.getLogger(__name__, 'reserve_disk_space')
class NotEnoughSpace(Exception):
def __init__(self, free, reserve, drive):
self.free = free
self.reserve = reserve
self.drive = drive
def reserve_disk_space(reserve, drive):
'''
Returns a dotdict containing these values:
{
'free': integer number of free bytes,
'reserve': integer number of requested bytes,
'drive': path string representing the drive we checked
}
Raises NotEnoughSpace if the amount of free disk space on `drive` is less
than `free`.
'''
drive = os.path.abspath(drive)
drive = os.path.splitdrive(drive)[0]
log.debug('Checking drive %s.', drive)
free = shutil.disk_usage(drive).free
if free < reserve:
raise NotEnoughSpace(free=free, reserve=reserve, drive=drive)
else:
return dotdict.DotDict(free=free, reserve=reserve, drive=drive)
def reserve_disk_space_argparse(args): def reserve_disk_space_argparse(args):
reserve = bytestring.parsebytes(args.reserve)
drive = pathclass.Path(args.drive)
try: try:
status = reserve_disk_space(reserve=bytestring.parsebytes(args.reserve), drive=args.drive) free = drive.assert_disk_space(reserve)
free = bytestring.bytestring(status.free) free = bytestring.bytestring(free)
reserve = bytestring.bytestring(status.reserve) reserve = bytestring.bytestring(reserve)
log.info('There is %s available out of %s.', free, reserve) log.info('There is %s available out of %s.', free, reserve)
return 0 return 0
except NotEnoughSpace as exc: except pathclass.NotEnoughSpace as exc:
free = bytestring.bytestring(exc.free) free = bytestring.bytestring(exc.free)
reserve = bytestring.bytestring(exc.reserve) reserve = bytestring.bytestring(exc.reserve)
log.fatal('Only %s available out of %s.', free, reserve) log.fatal('Only %s available out of %s.', free, reserve)