import functools
import shutil


def save_workspace(new_workspace):
    """Decorator to save a workspace to a new location.

        If `new_workspace` already exists on disk, it will be recursively
        removed.

        Example usage with a test case::

            import natcap.invest.testing

            @natcap.invest.testing.save_workspace('/path/to/workspace')
            def test_workspaces(self):
                model.execute(self.args)

        Note:
            + Target workspace folder must be saved to ``self.workspace_dir``
                This decorator is only designed to work with test functions
                from subclasses of ``unittest.TestCase`` such as
                ``natcap.invest.testing.GISTest``.

            + If ``new_workspace`` exists, it will be removed.
                So be careful where you save things.

        Args:
            new_workspace (string): a URI to the where the workspace should be
                copied.

        Returns:
            A composed test case function which will execute and then save your
            workspace to the specified location."""

    # item is the function being decorated
    def test_inner_func(item):

        # this decorator indicates that this innermost function is wrapping up
        # the function passed in as item.
        @functools.wraps(item)
        def test_and_remove_workspace(self, *args, **kwargs):
            # This inner function actually executes the test function and then
            # moves the workspace to the folder passed in by the user.
            item(self)

            # remove the contents of the old folder
            try:
                shutil.rmtree(new_workspace)
            except OSError:
                pass

            # copy the workspace to the target folder
            old_workspace = self.workspace_dir
            shutil.copytree(old_workspace, new_workspace)
        return test_and_remove_workspace
    return test_inner_func