| 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) |