Rough pythonisation of disktuil apfs -list -plist

apfs_disks.py

#!/usr/bin/python
'''List APFS containers and return properties. Free to use/abuse/improve/laugh at'''

import plistlib
import subprocess
import sys

from collections import namedtuple


# Create the namedtuple collections
# General information about the container
Container = namedtuple('Container', ['APFSContainerUUID', 'CapacityCeiling', 'CapacityFree',  # NOQA
                                     'ContainerReference', 'DesignatedPhysicalStore', 'Fusion',  # NOQA
                                     'PhysicalStores', 'Volumes'])
# There is a PhysicalStores for each container.
# There can be more than one physical store in a container
PhysicalStores = namedtuple('PhysicalStores', ['DeviceIdentifier', 'DiskUUID', 'Size'])  # NOQA

# Volume information (roughly equates to partition).
# There can be more than one volume per container
Volume = namedtuple('Volume', ['APFSVolumeUUID', 'CapacityInUse', 'CapacityQuota',  # NOQA
                               'CapacityReserve', 'CryptoMigrationOn', 'DeviceIdentifier',  # NOQA
                               'Encryption', 'Locked', 'Name', 'Roles'])  # NOQA


# diskutil apfs list output as plist
def get_apfs_list():
    cmd = ['/usr/sbin/diskutil', 'apfs', 'list', '-plist']
    (result, error) = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE).communicate()  # NOQA

    if result:
        diskutil = plistlib.readPlistFromString(result.strip('\n'))
        return diskutil
    elif error:
        log(error)
        sys.exit(1)


def disk_generator(volume_name=None):
    '''Returns a generator containing namedtuple for all containers.'''
    # Create some easy to reference variables to use in processing.
    # disks = get_apfs_list()
    apfs_containers = get_apfs_list()['Containers']

    # Because we're really getting meta on namedtuples in namedtuples, use
    # functions to get the info required.
    def get_physical_store_info(container):
        '''Returns info about a physical store in relation to the specified
        container'''
        for physical_store in container['PhysicalStores']:
                    _physical_store = PhysicalStores(
                        DeviceIdentifier=physical_store['DeviceIdentifier'],
                        DiskUUID=physical_store['DiskUUID'],
                        Size=physical_store['Size']
                    )
                    yield _physical_store

    def get_volume_info(volumes):
        '''Returns info about a volume from specified volumes.'''
        for volume in volumes:
            _volume = Volume(
                APFSVolumeUUID=volume['APFSVolumeUUID'],
                CapacityInUse=volume['CapacityInUse'],
                CapacityQuota=volume['CapacityQuota'],
                CapacityReserve=volume['CapacityReserve'],
                CryptoMigrationOn=volume['CryptoMigrationOn'],
                DeviceIdentifier=volume['DeviceIdentifier'],
                Encryption=volume['Encryption'],
                Locked=volume['Locked'],
                Name=volume['Name'],
                Roles=volume['Roles']
            )
            if volume_name and volume_name in _volume.Name:
                # Because return can't exist in a generator,
                # break after match is yielded.
                yield _volume
                break
            else:
                yield _volume

    # Loop through the containers list
    for container in apfs_containers:
        # Return a list if more than one, otherwise, don't return list.
        volumes = [x for x in get_volume_info(container['Volumes'])]
        if len(volumes) == 1:
            volumes = volumes[0]

        _container = Container(
            APFSContainerUUID=container['APFSContainerUUID'],
            CapacityCeiling=container['CapacityCeiling'],
            CapacityFree=container['CapacityFree'],
            ContainerReference=container['ContainerReference'],
            DesignatedPhysicalStore=container['DesignatedPhysicalStore'],
            Fusion=container['Fusion'],
            PhysicalStores=[x for x in get_physical_store_info(container)],
            Volumes=volumes
        )
        yield _container


def disk(volume_name):
    '''Returns a namedtuple for a single volume containing the
    Container, Physical Store Info and Volume.
    For example, to get the Container Reference disk for "Macintosh HD", use:
        disk('Macintosh HD').ContainerReference
    To get the APFSVolumeID for "Macintosh HD":
        disk('Macintosh HD').Volumes.APFSVolumeUUID'''
    disks = [x for x in disk_generator(volume_name)]
    return disks[0]