mysos/common/pkgutil.py

import os

from twitter.common.dirutil import safe_mkdir

import pkg_resources


def unpack_assets(output_dir, module, asset_path, execute=lambda x: None):
  """
    Extract files from a module in a package into 'output_dir'.

    :param output_dir: The directory to copy the assets *root* to.
    :param module: The module in the package to find the assets in. e.g., mysos.executor.
    :param asset_path: The path of the asset relative to the module root to unpack.
    :param execute: A function that expects a single argument as the path to an asset file. It can
                    be used to process this file.

    NOTE: If the specified 'asset_path' is a directory, its contents are copied but it itself is not
          recreated in the output directory. e.g., <module_dir>/files/bin/file.sh is copied to
          <output_dir>/bin/file.sh. An analogy is that its like `cp -R asset_path/* output_dir` and
          not like `cp -R asset_path output_dir`.
  """
  _unpack_assets(output_dir, module, asset_path, execute, asset_path)


def _unpack_assets(output_dir, module, asset_root, execute, current_path):
  """
    The internal helper function for unpack_assets(...) recursion.
    :param current_path: Records the current
  """
  for asset in pkg_resources.resource_listdir(module, current_path):
    asset_target = os.path.join(os.path.relpath(current_path, asset_root), asset)
    if pkg_resources.resource_isdir(module, os.path.join(current_path, asset)):
      safe_mkdir(os.path.join(output_dir, asset_target))
      _unpack_assets(output_dir, module, asset_root, execute, os.path.join(current_path, asset))
    else:
      output_file = os.path.join(output_dir, asset_target)
      with open(output_file, 'wb') as fp:
        fp.write(pkg_resources.resource_string(
            module, os.path.join(asset_root, asset_target)))
        execute(output_file)