environment: Support options instead of properties

Support passing in an options struct instead of just using module
properties. Also take steps to remove all state from this module.

Bug: 600
Change-Id: Ief7be56e1b284e852a9b882977ec18d60dc8f776
Reviewed-on: https://pigweed-review.googlesource.com/c/infra/recipes/+/79900
Reviewed-by: Oliver Newman <olivernewman@google.com>
Commit-Queue: Rob Mohr <mohrr@google.com>
diff --git a/recipe_modules/environment/api.py b/recipe_modules/environment/api.py
index 64513b6..3c2f36b 100644
--- a/recipe_modules/environment/api.py
+++ b/recipe_modules/environment/api.py
@@ -14,22 +14,44 @@
 """Environment utility functions.
 
 Usage:
-api.environment.initialize(checkout_root=...)
-with api.environment():
+env = api.environment.initialize(checkout_root, env_options)
+with env():
   ...
 """
 
 import contextlib
 import pprint
 
+from PB.recipe_modules.pigweed.environment.options import Options
+
 import attr
 from recipe_engine import recipe_api
 
 
 @attr.s
-class Package(object):
-    name = attr.ib(type=bytes)
-    version = attr.ib(type=bytes)
+class Environment(object):
+    _api = attr.ib()
+    prefixes = attr.ib(default=attr.Factory(dict))
+    suffixes = attr.ib(default=attr.Factory(dict))
+    env = attr.ib(default=attr.Factory(dict))
+
+    @contextlib.contextmanager
+    def __call__(self):
+        # Using reversed() because things that are added later in environment
+        # setup need to override things that came earlier.
+        with self._api.context(
+            env_prefixes={k: reversed(v) for k, v in self.prefixes.items()},
+            env_suffixes=self.suffixes,
+            env=self.env,
+        ):
+            with self._api.macos_sdk():
+                yield self
+
+    def __getattr__(self, name):
+        if name not in self.env:
+            raise AttributeError(name)
+
+        return self.env.get(name)
 
 
 class EnvironmentApi(recipe_api.RecipeApi):
@@ -37,41 +59,39 @@
 
     def __init__(self, props, *args, **kwargs):
         super(EnvironmentApi, self).__init__(*args, **kwargs)
-        self._root_variable_name = str(props.root_variable_name)
-        self._relative_pigweed_root = str(props.relative_pigweed_root)
-        self._config_file = str(props.config_file)
-        self._cipd_dir = None
-        self._prefixes = {}
-        self._suffixes = {}
-        self._env = {}
-        self._initialized = False
-        self._cipd_installation_dirs = []
+        # TODO(pwbug/600) Remove self._options.
+        self._options = Options()
+        self._options.root_variable_name = str(props.root_variable_name)
+        self._options.relative_pigweed_root = str(props.relative_pigweed_root)
+        self._options.config_file = str(props.config_file)
+        # TODO(pwbug/600) Remove self._env.
+        self._env = None
 
-    def _init_platform(self):
+    def _init_platform(self, env):
         if self.m.platform.is_mac:
             with self.m.step.nest('setup platform'):
                 with self.m.macos_sdk():
                     pass
 
-    def _init_misc_vars(self):
-        self._env['PW_ENVIRONMENT_NO_ERROR_ON_UNRECOGNIZED'] = '1'
-        self._env['PW_ENVSETUP_DISABLE_SPINNER'] = '1'
-        self._env['PW_PRESUBMIT_DISABLE_SUBPROCESS_CAPTURE'] = '1'
+    def _init_misc_vars(self, env):
+        env.env['PW_ENVIRONMENT_NO_ERROR_ON_UNRECOGNIZED'] = '1'
+        env.env['PW_ENVSETUP_DISABLE_SPINNER'] = '1'
+        env.env['PW_PRESUBMIT_DISABLE_SUBPROCESS_CAPTURE'] = '1'
 
         if self.m.led.launched_by_led:
             # Not using self.m.buildbucket_util.id because having relatively
             # consistent length is important to some downstream projects when
             # pulling this into their builds.
-            self._env['BUILDBUCKET_ID'] = 'test:{}'.format(
+            env.env['BUILDBUCKET_ID'] = 'test:{}'.format(
                 self.m.swarming.task_id
             )
-            self._env['BUILD_NUMBER'] = str(self.m.buildbucket.build.number)
+            env.env['BUILD_NUMBER'] = str(self.m.buildbucket.build.number)
 
         else:
-            self._env['BUILDBUCKET_ID'] = str(self.m.buildbucket.build.id)
-            self._env['BUILD_NUMBER'] = str(self.m.buildbucket.build.number)
+            env.env['BUILDBUCKET_ID'] = str(self.m.buildbucket.build.id)
+            env.env['BUILD_NUMBER'] = str(self.m.buildbucket.build.number)
 
-        self._env['BUILDBUCKET_NAME'] = ':'.join(
+        env.env['BUILDBUCKET_NAME'] = ':'.join(
             (
                 self.m.buildbucket.build.builder.project,
                 self.m.buildbucket.build.builder.bucket,
@@ -79,20 +99,25 @@
             )
         )
 
-        if self._env['BUILDBUCKET_NAME'] == '::':
-            self._env['BUILDBUCKET_NAME'] = 'project:bucket:builder'
+        if env.env['BUILDBUCKET_NAME'] == '::':
+            env.env['BUILDBUCKET_NAME'] = 'project:bucket:builder'
 
-        self._env['GOCACHE'] = self.m.path['cache'].join('go')
-        self._env['PIP_CACHE_DIR'] = self.m.path['cache'].join('pip')
+        env.env['GOCACHE'] = self.m.path['cache'].join('go')
+        env.env['PIP_CACHE_DIR'] = self.m.path['cache'].join('pip')
         # Bazel cache dir.
-        self._env['TEST_TMPDIR'] = self.m.path['cache'].join('bazel')
+        env.env['TEST_TMPDIR'] = self.m.path['cache'].join('bazel')
 
-    def _init_pigweed(self, checkout_root, top_pres, use_constraint_file):
+    def _init_pigweed(
+        self,
+        checkout_root,
+        top_presentation,
+        use_constraint_file,
+        pigweed_root,
+        config_file,
+        env,
+    ):
         """Run pw_env_setup."""
 
-        if not self._config_file:
-            return
-
         def path(relative_path):
             parts = [
                 x for x in relative_path.split('/') if x not in ('.', u'.')
@@ -102,8 +127,6 @@
             else:
                 return checkout_root  # pragma: no cover
 
-        pw_root = path(self._relative_pigweed_root)
-
         env_dir = self.m.path['start_dir'].join('environment')
         json_file = env_dir.join('vars.json')
         shell_file = env_dir.join('setup.sh')
@@ -119,9 +142,11 @@
 
         cmd = [
             'python3',
-            pw_root.join('pw_env_setup', 'py', 'pw_env_setup', 'env_setup.py'),
+            pigweed_root.join(
+                'pw_env_setup', 'py', 'pw_env_setup', 'env_setup.py'
+            ),
             '--pw-root',
-            pw_root,
+            pigweed_root,
             '--install-dir',
             env_dir,
             '--json-file',
@@ -139,15 +164,15 @@
         if not use_constraint_file:
             cmd.append('--unpin-pip-packages')
 
-        cmd.extend(('--config-file', path(self._config_file)))
+        cmd.extend(('--config-file', path(config_file)))
 
-        top_pres.logs['vars.json'] = pprint.pformat(
-            self.m.file.read_json('read config', path(self._config_file))
+        top_presentation.logs['vars.json'] = pprint.pformat(
+            self.m.file.read_json('read config', path(config_file))
         )
 
         with self.m.step.defer_results():
             with self.m.step.nest('run pw_env_setup'):
-                with self():
+                with env():
                     self.m.step('pw_env_setup', cmd)
 
                 cipd_dir = env_dir.join('cipd')
@@ -185,56 +210,67 @@
             },
         )
 
-        top_pres.logs['vars.json'] = pprint.pformat(json_data)
+        top_presentation.logs['vars.json'] = pprint.pformat(json_data)
 
         for var, value in json_data['set'].items():
-            self._env[var] = value
+            env.env[var] = value
 
         for var, actions in json_data['modify'].items():
             for value in actions.get('prepend', ()):
-                self._prefixes.setdefault(var, [])
-                self._prefixes[var].append(value)
+                env.prefixes.setdefault(var, [])
+                env.prefixes[var].append(value)
             for value in actions.get('append', ()):
-                self._suffixes.setdefault(var, [])
-                self._suffixes[var].append(value)
+                env.suffixes.setdefault(var, [])
+                env.suffixes[var].append(value)
 
-    def init(self, checkout_root, use_constraint_file=True):
+    def init(
+        self, checkout_root, options=None, use_constraint_file=True,
+    ):
         pigweed_root = checkout_root
-        if self._relative_pigweed_root not in (None, '', '.'):
-            pigweed_root = checkout_root.join(self._relative_pigweed_root)
+        env = Environment(api=self.m)
 
-        if self._root_variable_name:
-            self._env[self._root_variable_name] = checkout_root
+        if not options:
+            options = self._options  # pragma: no cover
 
-        if not self._initialized:
-            with self.m.step.nest('environment') as pres:
-                # Setting _initialized immediately because some setup steps need
-                # to use the context of previous steps, and invoking self() is
-                # the easiest way to do so.
-                self._initialized = True
-                self._env['PW_ROOT'] = pigweed_root
-                self._env['PW_PROJECT_ROOT'] = checkout_root
+        if not options.config_file:
+            self._env = env
+            return env
 
-                self._init_platform()
-                self._init_misc_vars()
-                self._init_pigweed(checkout_root, pres, use_constraint_file)
+        if options.relative_pigweed_root not in (None, '', '.'):
+            pigweed_root = checkout_root.join(options.relative_pigweed_root)
+
+        if options.root_variable_name:
+            env.env[options.root_variable_name] = checkout_root
+
+        with self.m.step.nest('environment') as pres:
+            # Setting _initialized immediately because some setup steps need
+            # to use the context of previous steps, and invoking self() is
+            # the easiest way to do so.
+            env.env['PW_ROOT'] = pigweed_root
+            env.env['PW_PROJECT_ROOT'] = checkout_root
+
+            self._init_platform(env)
+            self._init_misc_vars(env)
+            self._init_pigweed(
+                checkout_root=checkout_root,
+                top_presentation=pres,
+                use_constraint_file=use_constraint_file,
+                pigweed_root=pigweed_root,
+                config_file=options.config_file,
+                env=env,
+            )
+
+        # TODO(pwbug/600) Remove self._env.
+        self._env = env
+        return env
 
     @contextlib.contextmanager
     def __call__(self):
-        assert self._initialized
-
-        # Using reversed() because things that are added later in environment
-        # setup need to override things that came earlier.
-        with self.m.context(
-            env_prefixes={k: reversed(v) for k, v in self._prefixes.items()},
-            env_suffixes=self._suffixes,
-            env=self._env,
-        ):
-            with self.m.macos_sdk():
-                yield self
+        # TODO(pwbug/600) Remove this method.
+        assert self._env
+        with self._env():
+            yield self
 
     def __getattr__(self, name):
-        if name not in self._env:
-            raise AttributeError(name)
-
-        return self._env.get(name)
+        assert self._env
+        return getattr(self._env, name)