1# Copyright 2022 The Pigweed Authors 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); you may not 4# use this file except in compliance with the License. You may obtain a copy of 5# the License at 6# 7# https://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT 11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 12# License for the specific language governing permissions and limitations under 13# the License. 14# pylint: disable=line-too-long 15"""Pigweed shell activation script. 16 17Aside from importing it, this script can be used in three ways: 18 191. Activate the Pigweed environment in your current shell (i.e., modify your 20 interactive shell's environment with Pigweed environment variables). 21 22 Using bash (assuming a global Python 3 is in $PATH): 23 source <(python3 ./pw_ide/activate.py -s bash) 24 25 Using bash (using the environment Python): 26 source <({environment}/pigweed-venv/bin/python ./pw_ide/activate.py -s bash) 27 282. Run a shell command or executable in an activated shell (i.e. apply a 29 modified environment to a subprocess without affecting your current 30 interactive shell). 31 32 Example (assuming a global Python 3 is in $PATH): 33 python3 ./pw_ide/activate.py -x 'pw ide cpp --list' 34 35 Example (using the environment Python): 36 {environment}/pigweed-venv/bin/python ./pw_ide/activate.py -x 'pw ide cpp --list' 37 38 Example (using the environment Python on Windows): 39 {environment}/pigweed-venv/Scripts/pythonw.exe ./pw_ide/activate.py -x 'pw ide cpp --list' 40 413. Produce a JSON representation of the Pigweed activated environment (-O) or 42 the diff against your current environment that produces an activated 43 environment (-o). See the help for more detailed information on the options 44 available. 45 46 Example (assuming a global Python 3 is in $PATH): 47 python3 ./pw_ide/activate.py -o 48 49 Example (using the environment Python): 50 {environment}/pigweed-venv/bin/python ./pw_ide/activate.py -o 51 52 Example (using the environment Python on Windows): 53 {environment}/pigweed-venv/Scripts/pythonw.exe ./pw_ide/activate.py -o 54""" 55# pylint: enable=line-too-long 56 57from __future__ import annotations 58 59from abc import abstractmethod, ABC 60import argparse 61from collections import defaultdict 62from inspect import cleandoc 63from functools import cache 64import json 65import os 66from pathlib import Path 67import shlex 68import subprocess 69import sys 70from typing import cast, Tuple 71 72 73def find_pigweed_json_above(working_dir=Path(os.getcwd())) -> Path | None: 74 """Find the path to pigweed.json by searching this directory and above. 75 76 This starts looking in the current working directory, then recursively in 77 each directory above the current working directory, until it finds a 78 pigweed.json file or reaches the file system root. So invoking this anywhere 79 within a Pigweed project directory should work. 80 """ 81 82 if (pigweed_json := (working_dir / "pigweed.json")).exists(): 83 return pigweed_json.parent 84 85 # Recursively search in directories above this one. 86 # This condition will be false when we reach the root of the file system. 87 if working_dir.parent != working_dir: 88 return find_pigweed_json_above(working_dir.parent) 89 90 return None 91 92 93def find_pigweed_json_below(working_dir=Path(os.getcwd())) -> Path | None: 94 """Find the path to pigweed.json by searching below this directory. 95 96 This will return the one nearest subdirectory that contains a pigweeed.json 97 file. 98 """ 99 100 # We only want the first result (there could be many other spurious 101 # pigweed.json files in environment directories), but we can't directly 102 # index a generator, and getting a list by running the generator to 103 # completion is pointlessly slow. 104 for path in working_dir.rglob('pigweed.json'): 105 dir_path = path.parent 106 # There's a source file for pigweed.json that we want to omit. 107 if dir_path.parent.name != "cipd_setup": 108 return dir_path 109 110 return None 111 112 113@cache 114def pigweed_root( 115 working_dir=Path(os.getcwd()), use_env_vars=True 116) -> Tuple[Path, bool]: 117 """Find the Pigweed root within the project. 118 119 The presence of a pigweed.json file is the sentinel for the Pigweed root. 120 The heuristic is to first search in the current directory or above ("are 121 we inside of a Pigweed directory?"), and failing that, to search in the 122 directories below ("does this project contain a Pigweed directory?"). 123 124 This returns two values: the Pigweed root directory, and a boolean 125 indicating that the directory path is "validated". A validated path came 126 from a source of truth, like the location of a pigweed.json file or an 127 environment variable. An unvalidated path may just be a last resort guess. 128 129 Note that this logic presumes that there's only one Pigweed project 130 directory. In a hypothetical project setup that contained multiple Pigweed 131 projects, this would continue to work when invoked inside of one of those 132 Pigweed directories, but would have inconsistent results when invoked 133 in a parent directory. 134 """ 135 # These are shortcuts that make this faster if we happen to be running in 136 # an activated environment. This can be disabled, e.g., in tests. 137 if use_env_vars: 138 if (pw_root := os.environ.get('PW_ROOT')) is not None: 139 return Path(pw_root), True 140 141 if (pw_project_root := os.environ.get('PW_PROJECT_ROOT')) is not None: 142 return Path(pw_project_root), True 143 144 # If we're not in an activated environment (which is typical for this 145 # module), search for the pigweed.json sentinel. 146 if (root_above := find_pigweed_json_above(working_dir)) is not None: 147 return root_above, True 148 149 if (root_below := find_pigweed_json_below(working_dir)) is not None: 150 return root_below, True 151 152 return Path(os.getcwd()), False 153 154 155@cache 156def assumed_environment_root() -> Path | None: 157 """Infer the path to the Pigweed environment directory. 158 159 First we look at the environment variable that should contain the path if 160 we're operating in an activated Pigweed environment. If we don't find the 161 path there, we check a few known locations. If we don't find an environment 162 directory in any of those locations, we return None. 163 """ 164 actual_environment_root = os.environ.get('_PW_ACTUAL_ENVIRONMENT_ROOT') 165 166 if ( 167 actual_environment_root is not None 168 and (root_path := Path(actual_environment_root)).exists() 169 ): 170 return root_path.absolute() 171 172 root, root_is_validated = pigweed_root() 173 174 if not root_is_validated: 175 return None 176 177 default_environment = root / 'environment' 178 if default_environment.exists(): 179 return default_environment.absolute() 180 181 default_dot_environment = root / '.environment' 182 if default_dot_environment.exists(): 183 return default_dot_environment.absolute() 184 185 return None 186 187 188# We're looking for the `actions.json` file that allows us to activate the 189# Pigweed environment. That file is located in the Pigweed environment 190# directory, so if we found an environment directory, this variable will 191# have the path to `actions.json`. If it doesn't find an environment directory 192# (e.g., this isn't being executed in the context of a Pigweed project), this 193# will be None. Note that this is the "default" config file path because 194# callers of functions that need this path can provide their own paths to an 195# `actions.json` file. 196_DEFAULT_CONFIG_FILE_PATH = ( 197 None 198 if assumed_environment_root() is None 199 else cast(Path, assumed_environment_root()) / 'actions.json' 200) 201 202 203def _sanitize_path( 204 path: str, project_root_prefix: str, user_home_prefix: str 205) -> str: 206 """Given a path, return a sanitized path. 207 208 By default, environment variable paths are usually absolute. If we want 209 those paths to work across multiple systems, we need to sanitize them. This 210 takes a string that may be a path, and if it is indeed a path, it returns 211 the sanitized path, which is relative to either the repository root or the 212 user's home directory. If it's not a path, it just returns the input. 213 214 You can provide the strings that should be substituted for the project root 215 and the user's home directory. This may be useful for applications that have 216 their own way of representing those directories. 217 218 Note that this is intended to work on Pigweed environment variables, which 219 should all be relative to either of those two locations. Paths that aren't 220 (e.g. the path to a system binary) won't really be sanitized. 221 """ 222 # Return the argument if it's not actually a path. 223 # This strategy relies on the fact that env_setup outputs absolute paths for 224 # all path env vars. So if we get a variable that's not an absolute path, it 225 # must not be a path at all. 226 if not Path(path).is_absolute(): 227 return path 228 229 root, _ = pigweed_root() 230 231 project_root = root.resolve() 232 user_home = Path.home().resolve() 233 resolved_path = Path(path).resolve() 234 235 if resolved_path.is_relative_to(project_root): 236 return f'{project_root_prefix}/' + str( 237 resolved_path.relative_to(project_root) 238 ) 239 240 if resolved_path.is_relative_to(user_home): 241 return f'{user_home_prefix}/' + str( 242 resolved_path.relative_to(user_home) 243 ) 244 245 # Path is not in the project root or user home, so just return it as is. 246 return path 247 248 249class ShellModifier(ABC): 250 """Abstract class for shell modifiers. 251 252 A shell modifier provides an interface for modifying the environment 253 variables in various shells. You can pass in a current environment state 254 as a dictionary during instantiation and modify it and/or modify shell state 255 through other side effects. 256 """ 257 258 separator = ':' 259 comment = '# ' 260 261 def __init__( 262 self, 263 env: dict[str, str] | None = None, 264 env_only: bool = False, 265 path_var: str = '$PATH', 266 project_root: str = '.', 267 user_home: str = '~', 268 ): 269 # This will contain only the modifications to the environment, with 270 # no elements of the existing environment aside from variables included 271 # here. In that sense, it's like a diff against the existing 272 # environment, or a structured form of the shell modification side 273 # effects. 274 default_env_mod = {'PATH': path_var} 275 self.env_mod = default_env_mod.copy() 276 277 # This is seeded with the existing environment, and then is modified. 278 # So it contains the complete new environment after modifications. 279 # If no existing environment is provided, this is identical to env_mod. 280 env = env if env is not None else default_env_mod.copy() 281 self.env: dict[str, str] = defaultdict(str, env) 282 283 # Will contain the side effects, i.e. commands executed in the shell to 284 # modify its environment. 285 self.side_effects = '' 286 287 # Set this to not do any side effects, but just modify the environment 288 # stored in this class. 289 self.env_only = env_only 290 291 self.project_root = project_root 292 self.user_home = user_home 293 294 def do_effect(self, effect: str): 295 """Add to the commands that will affect the shell's environment. 296 297 This is a no-op if the shell modifier is set to only store shell 298 modification data rather than doing the side effects. 299 """ 300 if not self.env_only: 301 self.side_effects += f'{effect}\n' 302 303 def modify_env( 304 self, 305 config_file_path: Path | None = _DEFAULT_CONFIG_FILE_PATH, 306 sanitize: bool = False, 307 ) -> ShellModifier: 308 """Modify the current shell state per the actions.json file provided.""" 309 json_file_options = {} 310 311 if config_file_path is None: 312 raise RuntimeError( 313 'This must be run from a bootstrapped Pigweed directory!' 314 ) 315 316 with config_file_path.open('r') as json_file: 317 json_file_options = json.loads(json_file.read()) 318 319 root = self.project_root 320 home = self.user_home 321 322 # Set env vars 323 for var_name, value in json_file_options.get('set', dict()).items(): 324 if value is not None: 325 value = _sanitize_path(value, root, home) if sanitize else value 326 self.set_variable(var_name, value) 327 328 # Prepend & append env vars 329 for var_name, mode_changes in json_file_options.get( 330 'modify', dict() 331 ).items(): 332 for mode_name, values in mode_changes.items(): 333 if mode_name in ['prepend', 'append']: 334 modify_variable = self.prepend_variable 335 336 if mode_name == 'append': 337 modify_variable = self.append_variable 338 339 for value in values: 340 value = ( 341 _sanitize_path(value, root, home) 342 if sanitize 343 else value 344 ) 345 modify_variable(var_name, value) 346 347 return self 348 349 @abstractmethod 350 def set_variable(self, var_name: str, value: str) -> None: 351 pass 352 353 @abstractmethod 354 def prepend_variable(self, var_name: str, value: str) -> None: 355 pass 356 357 @abstractmethod 358 def append_variable(self, var_name: str, value: str) -> None: 359 pass 360 361 362class BashShellModifier(ShellModifier): 363 """Shell modifier for bash.""" 364 365 def set_variable(self, var_name: str, value: str): 366 self.env[var_name] = value 367 self.env_mod[var_name] = value 368 quoted_value = shlex.quote(value) 369 self.do_effect(f'export {var_name}={quoted_value}') 370 371 def prepend_variable(self, var_name: str, value: str) -> None: 372 self.env[var_name] = f'{value}{self.separator}{self.env[var_name]}' 373 self.env_mod[ 374 var_name 375 ] = f'{value}{self.separator}{self.env_mod[var_name]}' 376 quoted_value = shlex.quote(value) 377 self.do_effect( 378 f'export {var_name}={quoted_value}{self.separator}${var_name}' 379 ) 380 381 def append_variable(self, var_name: str, value: str) -> None: 382 self.env[var_name] = f'{self.env[var_name]}{self.separator}{value}' 383 self.env_mod[ 384 var_name 385 ] = f'{self.env_mod[var_name]}{self.separator}{value}' 386 quoted_value = shlex.quote(value) 387 self.do_effect( 388 f'export {var_name}=${var_name}{self.separator}{quoted_value}' 389 ) 390 391 392def _build_argument_parser() -> argparse.ArgumentParser: 393 """Set up `argparse`.""" 394 doc = __doc__ 395 396 try: 397 env_root = assumed_environment_root() 398 except RuntimeError: 399 env_root = None 400 401 # Substitute in the actual environment path in the help text, if we can 402 # find it. If not, leave the placeholder text. 403 if env_root is not None: 404 doc = doc.replace( 405 '{environment}', str(env_root.relative_to(Path.cwd())) 406 ) 407 408 parser = argparse.ArgumentParser( 409 formatter_class=argparse.RawDescriptionHelpFormatter, 410 description=doc, 411 ) 412 413 default_config_file_path = None 414 415 if _DEFAULT_CONFIG_FILE_PATH is not None: 416 default_config_file_path = _DEFAULT_CONFIG_FILE_PATH.relative_to( 417 Path.cwd() 418 ) 419 420 parser.add_argument( 421 '-c', 422 '--config-file', 423 default=_DEFAULT_CONFIG_FILE_PATH, 424 type=Path, 425 help='Path to actions.json config file, which defines ' 426 'the modifications to the shell environment ' 427 'needed to activate Pigweed. ' 428 f'Default: {default_config_file_path}', 429 ) 430 431 default_shell = Path(os.environ['SHELL']).name 432 parser.add_argument( 433 '-s', 434 '--shell-mode', 435 default=default_shell, 436 help='Which shell is being used. ' f'Default: {default_shell}', 437 ) 438 439 parser.add_argument( 440 '-o', 441 '--out', 442 action='store_true', 443 help='Write only the modifications to the environment ' 'out to JSON.', 444 ) 445 446 parser.add_argument( 447 '-O', 448 '--out-all', 449 action='store_true', 450 help='Write the complete modified environment to ' 'JSON.', 451 ) 452 453 parser.add_argument( 454 '-n', 455 '--sanitize', 456 action='store_true', 457 help='Sanitize paths that are relative to the repo ' 458 'root or user home directory so that they are portable ' 459 'to other workstations.', 460 ) 461 462 parser.add_argument( 463 '--path-var', 464 default='$PATH', 465 help='The string to substitute for the existing $PATH. Default: $PATH', 466 ) 467 468 parser.add_argument( 469 '--project-root', 470 default='.', 471 help='The string to substitute for the project root when sanitizing ' 472 'paths. Default: .', 473 ) 474 475 parser.add_argument( 476 '--user-home', 477 default='~', 478 help='The string to substitute for the user\'s home when sanitizing ' 479 'paths. Default: ~', 480 ) 481 482 parser.add_argument( 483 '-x', 484 '--exec', 485 help='A command to execute in the activated shell.', 486 metavar='COMMAND', 487 ) 488 489 return parser 490 491 492def main() -> int: 493 """The main CLI script.""" 494 args, _unused_extra_args = _build_argument_parser().parse_known_args() 495 env = os.environ.copy() 496 config_file_path = args.config_file 497 498 if not config_file_path.exists(): 499 sys.stderr.write(f'File not found! {config_file_path}') 500 sys.stderr.write( 501 'This must be run from a bootstrapped Pigweed ' 'project directory.' 502 ) 503 sys.exit(1) 504 505 # If we're executing a command in a subprocess, don't modify the current 506 # shell's state. Instead, apply the modified state to the subprocess. 507 env_only = args.exec is not None 508 509 # Assume bash by default. 510 shell_modifier = BashShellModifier 511 512 # TODO(chadnorvell): if args.shell_mode == 'zsh', 'ksh', 'fish'... 513 try: 514 modified_env = shell_modifier( 515 env=env, 516 env_only=env_only, 517 path_var=args.path_var, 518 project_root=args.project_root, 519 user_home=args.user_home, 520 ).modify_env(config_file_path, args.sanitize) 521 except (FileNotFoundError, json.JSONDecodeError): 522 sys.stderr.write( 523 'Unable to read file: {}\n' 524 'Please run this in bash or zsh:\n' 525 ' . ./bootstrap.sh\n'.format(str(config_file_path)) 526 ) 527 528 sys.exit(1) 529 530 if args.out_all: 531 print(json.dumps(modified_env.env, sort_keys=True, indent=2)) 532 return 0 533 534 if args.out: 535 print(json.dumps(modified_env.env_mod, sort_keys=True, indent=2)) 536 return 0 537 538 if args.exec is not None: 539 # Ensure that the command is always dequoted. 540 # When executed directly from the shell, this is already done by 541 # default. But in other contexts, the command may be passed more 542 # literally with whitespace and quotes, which won't work. 543 exec_cmd = args.exec.strip(" '") 544 545 # We're executing a command in a subprocess with the modified env. 546 return subprocess.run( 547 exec_cmd, env=modified_env.env, shell=True 548 ).returncode 549 550 # If we got here, we're trying to modify the current shell's env. 551 print(modified_env.side_effects) 552 553 # Let's warn the user if the output is going to stdout instead of being 554 # executed by the shell. 555 python_path = Path(sys.executable).relative_to(os.getcwd()) 556 c = shell_modifier.comment # pylint: disable=invalid-name 557 print( 558 cleandoc( 559 f""" 560 {c} 561 {c}Can you see these commands? If so, you probably wanted to 562 {c}source this script instead of running it. Try this instead: 563 {c} 564 {c} . <({str(python_path)} {' '.join(sys.argv)}) 565 {c} 566 {c}Run this script with `-h` for more help.""" 567 ) 568 ) 569 return 0 570 571 572if __name__ == '__main__': 573 sys.exit(main()) 574