├── .devcontainer ├── Dockerfile └── devcontainer.json ├── .editorconfig ├── .github ├── ISSUE_TEMPLATE │ ├── bug_report.md │ └── feature_request.md ├── PULL_REQUEST_TEMPLATE.md └── workflows │ ├── e2etests.yml │ └── pre-commit.yml ├── .gitignore ├── .isort.cfg ├── .pre-commit-config.yaml ├── .pylintrc ├── .style.yapf ├── .vscode └── settings.json ├── LICENSE ├── README.md ├── docker ├── Dockerfile ├── README.md ├── docker-build.yml ├── docker-compose.yml ├── jupyter_notebook_config.py └── notebook_init.py ├── docs ├── AddingAMagic.md ├── Contributing.md ├── DevelopmentEnvironment.md ├── Installation.md └── code-of-conduct.md ├── end_to_end_tests ├── __init__.py ├── basic_test.py ├── interface.py ├── manager.py ├── timesketch.py └── tools │ └── run_in_container.py ├── install.sh ├── notebooks ├── Quick_Primer_on_Colab_Jupyter.ipynb └── adding_magic.ipynb ├── picatrix ├── __init__.py ├── helpers │ ├── __init__.py │ └── table.py ├── lib │ ├── __init__.py │ ├── error.py │ ├── framework.py │ ├── framework_test.py │ ├── ipython.py │ ├── magic.py │ ├── magic_test.py │ ├── manager.py │ ├── manager_test.py │ ├── namespace.py │ ├── namespace_test.py │ ├── state.py │ ├── state_test.py │ ├── testlib.py │ ├── utils.py │ └── utils_test.py ├── magics │ ├── __init__.py │ ├── common.py │ ├── common_test.py │ └── timesketch.py ├── notebook_init.py └── version.py ├── prepare-picatrix-runtime.sh ├── requirements.txt ├── requirements_dev.txt ├── requirements_runtime.txt └── setup.py /.devcontainer/Dockerfile: -------------------------------------------------------------------------------- 1 | # See here for image contents: https://github.com/microsoft/vscode-dev-containers/tree/v0.202.3/containers/python-3/.devcontainer/base.Dockerfile 2 | 3 | # [Choice] Python version (use -bullseye variants on local arm64/Apple Silicon): 3, 3.9, 3.8, 3.7, 3.6, 3-bullseye, 3.9-bullseye, 3.8-bullseye, 3.7-bullseye, 3.6-bullseye, 3-buster, 3.9-buster, 3.8-buster, 3.7-buster, 3.6-buster 4 | ARG VARIANT="3.9-bullseye" 5 | FROM mcr.microsoft.com/vscode/devcontainers/python:0-${VARIANT} 6 | 7 | # [Choice] Node.js version: none, lts/*, 16, 14, 12, 10 8 | ARG NODE_VERSION="none" 9 | RUN if [ "${NODE_VERSION}" != "none" ]; then su vscode -c "umask 0002 && . /usr/local/share/nvm/nvm.sh && nvm install ${NODE_VERSION} 2>&1"; fi 10 | 11 | # [Optional] If your pip requirements rarely change, uncomment this section to add them to the image. 12 | # COPY requirements.txt /tmp/pip-tmp/ 13 | # RUN pip3 --disable-pip-version-check --no-cache-dir install -r /tmp/pip-tmp/requirements.txt \ 14 | # && rm -rf /tmp/pip-tmp 15 | 16 | # [Optional] Uncomment this section to install additional OS packages. 17 | # RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \ 18 | # && apt-get -y install --no-install-recommends 19 | 20 | # [Optional] Uncomment this line to install global node packages. 21 | # RUN su vscode -c "source /usr/local/share/nvm/nvm.sh && npm install -g " 2>&1 22 | -------------------------------------------------------------------------------- /.devcontainer/devcontainer.json: -------------------------------------------------------------------------------- 1 | // For format details, see https://aka.ms/devcontainer.json. For config options, see the README at: 2 | // https://github.com/microsoft/vscode-dev-containers/tree/v0.202.3/containers/python-3 3 | { 4 | "name": "Python 3", 5 | "runArgs": [ 6 | "--init" 7 | ], 8 | "build": { 9 | "dockerfile": "Dockerfile", 10 | "context": "..", 11 | "args": { 12 | // Update 'VARIANT' to pick a Python version: 3, 3.9, 3.8, 3.7, 3.6. 13 | // Append -bullseye or -buster to pin to an OS version. 14 | // Use -bullseye variants on local on arm64/Apple Silicon. 15 | "VARIANT": "3.7", 16 | // Options 17 | "NODE_VERSION": "lts/*" 18 | } 19 | }, 20 | // Set *default* container specific settings.json values on container create. 21 | "settings": { 22 | "python.pythonPath": "/usr/local/bin/python", 23 | "python.languageServer": "Pylance", 24 | "python.formatting.autopep8Path": "/usr/local/py-utils/bin/autopep8", 25 | "python.formatting.blackPath": "/usr/local/py-utils/bin/black", 26 | "python.formatting.yapfPath": "/usr/local/py-utils/bin/yapf", 27 | "python.linting.banditPath": "/usr/local/py-utils/bin/bandit", 28 | "python.linting.flake8Path": "/usr/local/py-utils/bin/flake8", 29 | "python.linting.mypyPath": "/usr/local/py-utils/bin/mypy", 30 | "python.linting.pycodestylePath": "/usr/local/py-utils/bin/pycodestyle", 31 | "python.linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle", 32 | "python.linting.pylintPath": "/usr/local/py-utils/bin/pylint", 33 | }, 34 | // Add the IDs of extensions you want installed when the container is created. 35 | "extensions": [ 36 | "ms-python.python", 37 | "ms-python.vscode-pylance" 38 | ], 39 | // Use 'forwardPorts' to make a list of ports inside the container available locally. 40 | "forwardPorts": [ 41 | 8888 42 | ], 43 | // Use 'postCreateCommand' to run commands after the container is created. 44 | "postCreateCommand": "pip3 install --user -r requirements_dev.txt; pip install -e .[runtime]; npm i pyright", 45 | // Comment out connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root. 46 | "remoteUser": "vscode" 47 | } 48 | -------------------------------------------------------------------------------- /.editorconfig: -------------------------------------------------------------------------------- 1 | root = true 2 | 3 | [*] 4 | charset = utf-8 5 | end_of_line = lf 6 | insert_final_newline = true 7 | trim_trailing_whitespace = true 8 | 9 | [*.py] 10 | indent_style = space 11 | indent_size = 2 12 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: Bug 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Describe the bug** 11 | A clear and concise description of what the bug is. 12 | 13 | **To Reproduce** 14 | Steps to reproduce the behavior: 15 | 1. Go to '...' 16 | 2. Click on '....' 17 | 3. Scroll down to '....' 18 | 4. See error 19 | 20 | **Expected behavior** 21 | A clear and concise description of what you expected to happen. 22 | 23 | **Screenshots** 24 | If applicable, add screenshots to help explain your problem. 25 | 26 | **Desktop (please complete the following information):** 27 | - OS: [e.g. iOS] 28 | - Browser [e.g. chrome, safari] 29 | - Version [e.g. 22] 30 | 31 | **Additional context** 32 | Add any other context about the problem here. 33 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature request 3 | about: Suggest an idea for this project 4 | title: '' 5 | labels: Feature request 6 | assignees: kiddinn, bladyjoker, mariuszlitwin 7 | 8 | --- 9 | 10 | **Describe the problem statement you are attempting to solve. Is the feature request related to 11 | a problem? Please describe.** 12 | A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] 13 | 14 | **Describe the solution you'd like** 15 | A clear and concise description of what you want to happen. 16 | 17 | **Describe alternatives you've considered** 18 | A clear and concise description of any alternative solutions or features you've considered. 19 | 20 | **Additional context** 21 | Add any other context or screenshots about the feature request here. 22 | -------------------------------------------------------------------------------- /.github/PULL_REQUEST_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | This PR fixes # 2 | 3 | > It's a good idea to open an issue first for discussion. 4 | > Describe in a sentence or few what the PR accomplishes. 5 | 6 | - [ ] Unit tests added 7 | - [ ] End-to-end tests added 8 | - [ ] Appropriate changes to documentation is included 9 | - [ ] If additional dependencies are needed, are they added into dependency files. 10 | -------------------------------------------------------------------------------- /.github/workflows/e2etests.yml: -------------------------------------------------------------------------------- 1 | name: picatrix-end-to-end 2 | on: 3 | pull_request: 4 | types: [opened, synchronize, reopened] 5 | schedule: 6 | - cron: '30 11 * * *' 7 | jobs: 8 | end2end: 9 | runs-on: ubuntu-20.04 10 | steps: 11 | - uses: actions/checkout@v2 12 | - name: Set up infrastructure with docker-compose 13 | run: docker-compose -f docker/docker-build.yml up -d 14 | env: 15 | JUPYTER_PORT: 8899 16 | - name: Run e2e tests 17 | run: docker-compose -f docker/docker-build.yml exec -e PYTHONPATH="." -w /usr/local/src/picatrix -T picatrix /home/picatrix/picenv/bin/python /usr/local/src/picatrix/end_to_end_tests/tools/run_in_container.py 18 | -------------------------------------------------------------------------------- /.github/workflows/pre-commit.yml: -------------------------------------------------------------------------------- 1 | name: pre-commit 2 | 3 | on: [push, pull_request] 4 | 5 | env: 6 | LANG: "C.UTF-8" 7 | LC_ALL: "C.UTF-8" 8 | jobs: 9 | pre-commit: 10 | runs-on: ubuntu-20.04 11 | strategy: 12 | matrix: 13 | os: [ubuntu-20.04, macos-latest] 14 | python-version: [3.7, 3.8] 15 | 16 | steps: 17 | - uses: actions/checkout@v2 18 | - name: Set up Python ${{ matrix.python-version }} 19 | uses: actions/setup-python@v2 20 | with: 21 | python-version: ${{ matrix.python-version }} 22 | - name: Install dependencies 23 | run: | 24 | python setup.py install 25 | pip install -r requirements_dev.txt 26 | pip install pre-commit 27 | - name: Run all pre-commit checks 28 | run: | 29 | pre-commit run --all-files 30 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | 2 | # Created by https://www.toptal.com/developers/gitignore/api/python,jupyternotebooks 3 | # Edit at https://www.toptal.com/developers/gitignore?templates=python,jupyternotebooks 4 | 5 | ### JupyterNotebooks ### 6 | # gitignore template for Jupyter Notebooks 7 | # website: http://jupyter.org/ 8 | 9 | .ipynb_checkpoints 10 | */.ipynb_checkpoints/* 11 | 12 | # IPython 13 | profile_default/ 14 | ipython_config.py 15 | 16 | # Remove previous ipynb_checkpoints 17 | # git rm -r .ipynb_checkpoints/ 18 | 19 | ### Python ### 20 | # Byte-compiled / optimized / DLL files 21 | __pycache__/ 22 | *.py[cod] 23 | *$py.class 24 | 25 | # C extensions 26 | *.so 27 | 28 | # Distribution / packaging 29 | .Python 30 | build/ 31 | develop-eggs/ 32 | dist/ 33 | downloads/ 34 | eggs/ 35 | .eggs/ 36 | # lib/ # ignored as `lib` is a valid path in Picatrix 37 | lib64/ 38 | parts/ 39 | sdist/ 40 | var/ 41 | wheels/ 42 | share/python-wheels/ 43 | *.egg-info/ 44 | .installed.cfg 45 | *.egg 46 | MANIFEST 47 | 48 | node_modules 49 | package-lock.json 50 | 51 | # PyInstaller 52 | # Usually these files are written by a python script from a template 53 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 54 | *.manifest 55 | *.spec 56 | 57 | # Installer logs 58 | pip-log.txt 59 | pip-delete-this-directory.txt 60 | 61 | # Unit test / coverage reports 62 | htmlcov/ 63 | .tox/ 64 | .nox/ 65 | .coverage 66 | .coverage.* 67 | .cache 68 | nosetests.xml 69 | coverage.xml 70 | *.cover 71 | *.py,cover 72 | .hypothesis/ 73 | .pytest_cache/ 74 | cover/ 75 | 76 | # Translations 77 | *.mo 78 | *.pot 79 | 80 | # Django stuff: 81 | *.log 82 | local_settings.py 83 | db.sqlite3 84 | db.sqlite3-journal 85 | 86 | # Flask stuff: 87 | instance/ 88 | .webassets-cache 89 | 90 | # Scrapy stuff: 91 | .scrapy 92 | 93 | # Sphinx documentation 94 | docs/_build/ 95 | 96 | # PyBuilder 97 | .pybuilder/ 98 | target/ 99 | 100 | # Jupyter Notebook 101 | 102 | # IPython 103 | 104 | # pyenv 105 | # For a library or package, you might want to ignore these files since the code is 106 | # intended to run in multiple environments; otherwise, check them in: 107 | # .python-version 108 | 109 | # pipenv 110 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 111 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 112 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 113 | # install all needed dependencies. 114 | #Pipfile.lock 115 | 116 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 117 | __pypackages__/ 118 | 119 | # Celery stuff 120 | celerybeat-schedule 121 | celerybeat.pid 122 | 123 | # SageMath parsed files 124 | *.sage.py 125 | 126 | # Environments 127 | .env 128 | .venv 129 | env/ 130 | venv/ 131 | ENV/ 132 | env.bak/ 133 | venv.bak/ 134 | 135 | # Spyder project settings 136 | .spyderproject 137 | .spyproject 138 | 139 | # Rope project settings 140 | .ropeproject 141 | 142 | # mkdocs documentation 143 | /site 144 | 145 | # mypy 146 | .mypy_cache/ 147 | .dmypy.json 148 | dmypy.json 149 | 150 | # Pyre type checker 151 | .pyre/ 152 | 153 | # pytype static type analyzer 154 | .pytype/ 155 | 156 | # Cython debug symbols 157 | cython_debug/ 158 | 159 | # End of https://www.toptal.com/developers/gitignore/api/python,jupyternotebooks 160 | -------------------------------------------------------------------------------- /.isort.cfg: -------------------------------------------------------------------------------- 1 | [isort] 2 | multi_line_output = 3 3 | include_trailing_comma = True 4 | force_grid_wrap = 0 5 | use_parentheses = True 6 | ensure_newline_before_comments = True 7 | line_length = 80 8 | known_third_party = IPython,dateutil,dfdatetime,docstring_parser,ipyaggrid,mock,pandas,pytest,setuptools,timesketch_api_client,timesketch_import_client 9 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pre-commit/pre-commit-hooks 3 | rev: v4.0.1 4 | hooks: 5 | - id: check-yaml 6 | - id: end-of-file-fixer 7 | - id: trailing-whitespace 8 | - repo: https://github.com/pre-commit/mirrors-yapf 9 | rev: "v0.31.0" 10 | hooks: 11 | - id: yapf 12 | args: ["--style=.style.yapf", "--parallel", "--in-place"] 13 | - repo: https://github.com/asottile/seed-isort-config 14 | rev: v2.2.0 15 | hooks: 16 | - id: seed-isort-config 17 | - repo: https://github.com/pycqa/isort 18 | rev: 5.9.3 19 | hooks: 20 | - id: isort 21 | name: isort (python) 22 | verbose: true 23 | - repo: local 24 | hooks: 25 | - id: pylint 26 | name: pylint 27 | entry: pylint 28 | args: [picatrix] 29 | language: system 30 | types: [python] 31 | # To be enabled later when all code type checks 32 | # - id: pyright 33 | # name: pyright 34 | # entry: pyright 35 | # language: node 36 | # types: [python] 37 | # pass_filenames: false 38 | # additional_dependencies: ["pyright@1.1.130"] 39 | - id: pytest-check 40 | name: pytest-check 41 | entry: pytest 42 | language: system 43 | pass_filenames: false 44 | always_run: true 45 | -------------------------------------------------------------------------------- /.pylintrc: -------------------------------------------------------------------------------- 1 | # Pylint 2.1.x - 2.2.x configuration file 2 | # 3 | # This file is generated by l2tdevtools update-dependencies.py, any dependency 4 | # related changes should be made in dependencies.ini. 5 | [MASTER] 6 | 7 | # A comma-separated list of package or module names from where C extensions may 8 | # be loaded. Extensions are loading into the active Python interpreter and may 9 | # run arbitrary code 10 | extension-pkg-whitelist= 11 | 12 | # Add files or directories to the blacklist. They should be base names, not 13 | # paths. 14 | ignore=CVS,jupyter_notebook_config.py 15 | 16 | # Add files or directories matching the regex patterns to the blacklist. The 17 | # regex matches against base names, not paths. 18 | ignore-patterns= 19 | 20 | # Python code to execute, usually for sys.path manipulation such as 21 | # pygtk.require(). 22 | #init-hook= 23 | 24 | # Use multiple processes to speed up Pylint. 25 | jobs=1 26 | 27 | # List of plugins (as comma separated values of python modules names) to load, 28 | # usually to register additional checkers. 29 | load-plugins=pylint.extensions.docparams 30 | 31 | # Pickle collected data for later comparisons. 32 | persistent=yes 33 | 34 | # Specify a configuration file. 35 | #rcfile= 36 | 37 | # Allow loading of arbitrary C extensions. Extensions are imported into the 38 | # active Python interpreter and may run arbitrary code. 39 | unsafe-load-any-extension=no 40 | 41 | 42 | [MESSAGES CONTROL] 43 | 44 | # Only show warnings with the listed confidence levels. Leave empty to show 45 | # all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE, UNDEFINED 46 | confidence= 47 | 48 | # Disable the message, report, category or checker with the given id(s). You 49 | # can either give multiple identifiers separated by comma (,) or put this 50 | # option multiple times (only on the command line, not in the configuration 51 | # file where it should appear only once).You can also use "--disable=all" to 52 | # disable everything first and then reenable specific checks. For example, if 53 | # you want to run only the similarities checker, you can use "--disable=all 54 | # --enable=similarities". If you want to run only the classes checker, but have 55 | # no Warning level messages displayed, use"--disable=all --enable=classes 56 | # --disable=W" 57 | # 58 | disable= 59 | super-with-arguments, 60 | assignment-from-none, 61 | bad-inline-option, 62 | deprecated-pragma, 63 | duplicate-code, 64 | eq-without-hash, 65 | file-ignored, 66 | fixme, 67 | locally-disabled, 68 | locally-enabled, 69 | logging-format-interpolation, 70 | metaclass-assignment, 71 | missing-param-doc, 72 | no-absolute-import, 73 | no-self-use, 74 | parameter-unpacking, 75 | raw-checker-failed, 76 | suppressed-message, 77 | too-few-public-methods, 78 | too-many-ancestors, 79 | too-many-boolean-expressions, 80 | too-many-branches, 81 | too-many-instance-attributes, 82 | too-many-lines, 83 | too-many-locals, 84 | too-many-nested-blocks, 85 | too-many-public-methods, 86 | too-many-return-statements, 87 | too-many-statements, 88 | unsubscriptable-object, 89 | useless-object-inheritance, 90 | useless-suppression, 91 | no-else-raise, 92 | no-else-return, 93 | useless-type-doc, 94 | consider-using-f-string 95 | 96 | # Enable the message, report, category or checker with the given id(s). You can 97 | # either give multiple identifier separated by comma (,) or put this option 98 | # multiple time (only on the command line, not in the configuration file where 99 | # it should appear only once). See also the "--disable" option for examples. 100 | enable= 101 | 102 | 103 | [REPORTS] 104 | 105 | # Python expression which should return a note less than 10 (10 is the highest 106 | # note). You have access to the variables errors warning, statement which 107 | # respectively contain the number of errors / warnings messages and the total 108 | # number of statements analyzed. This is used by the global evaluation report 109 | # (RP0004). 110 | evaluation=10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10) 111 | 112 | # Template used to display messages. This is a python new-style format string 113 | # used to format the message information. See doc for all details 114 | #msg-template= 115 | 116 | # Set the output format. Available formats are text, parseable, colorized, json 117 | # and msvs (visual studio).You can also give a reporter class, eg 118 | # mypackage.mymodule.MyReporterClass. 119 | output-format=text 120 | 121 | # Tells whether to display a full report or only the messages 122 | reports=no 123 | 124 | # Activate the evaluation score. 125 | # score=yes 126 | score=no 127 | 128 | 129 | [REFACTORING] 130 | 131 | # Maximum number of nested blocks for function / method body 132 | max-nested-blocks=5 133 | 134 | 135 | [VARIABLES] 136 | 137 | # List of additional names supposed to be defined in builtins. Remember that 138 | # you should avoid to define new builtins when possible. 139 | additional-builtins= 140 | 141 | # Tells whether unused global variables should be treated as a violation. 142 | allow-global-unused-variables=yes 143 | 144 | # List of strings which can identify a callback function by name. A callback 145 | # name must start or end with one of those strings. 146 | callbacks=cb_,_cb 147 | 148 | # A regular expression matching the name of dummy variables (i.e. expectedly 149 | # not used). 150 | dummy-variables-rgx=_+$|(_[a-zA-Z0-9_]*[a-zA-Z0-9]+?$)|dummy|^ignored_|^unused_ 151 | 152 | # Argument names that match this expression will be ignored. Default to name 153 | # with leading underscore 154 | ignored-argument-names=_.*|^ignored_|^unused_ 155 | 156 | # Tells whether we should check for unused import in __init__ files. 157 | init-import=no 158 | 159 | # List of qualified module names which can have objects that can redefine 160 | # builtins. 161 | redefining-builtins-modules=six.moves,future.builtins 162 | 163 | 164 | [TYPECHECK] 165 | 166 | # List of decorators that produce context managers, such as 167 | # contextlib.contextmanager. Add to this list to register other decorators that 168 | # produce valid context managers. 169 | contextmanager-decorators=contextlib.contextmanager 170 | 171 | # List of members which are set dynamically and missed by pylint inference 172 | # system, and so shouldn't trigger E1101 when accessed. Python regular 173 | # expressions are accepted. 174 | generated-members= 175 | 176 | # Tells whether missing members accessed in mixin class should be ignored. A 177 | # mixin class is detected if its name ends with "mixin" (case insensitive). 178 | ignore-mixin-members=yes 179 | 180 | # This flag controls whether pylint should warn about no-member and similar 181 | # checks whenever an opaque object is returned when inferring. The inference 182 | # can return multiple potential results while evaluating a Python object, but 183 | # some branches might not be evaluated, which results in partial inference. In 184 | # that case, it might be useful to still emit no-member and other checks for 185 | # the rest of the inferred objects. 186 | ignore-on-opaque-inference=yes 187 | 188 | # List of class names for which member attributes should not be checked (useful 189 | # for classes with dynamically set attributes). This supports the use of 190 | # qualified names. 191 | ignored-classes=optparse.Values,thread._local,_thread._local 192 | 193 | # List of module names for which member attributes should not be checked 194 | # (useful for modules/projects where namespaces are manipulated during runtime 195 | # and thus existing member attributes cannot be deduced by static analysis. It 196 | # supports qualified module names, as well as Unix pattern matching. 197 | ignored-modules=grr_response_proto.* 198 | 199 | # Show a hint with possible names when a member name was not found. The aspect 200 | # of finding the hint is based on edit distance. 201 | missing-member-hint=yes 202 | 203 | # The minimum edit distance a name should have in order to be considered a 204 | # similar match for a missing member name. 205 | missing-member-hint-distance=1 206 | 207 | # The total number of similar names that should be taken in consideration when 208 | # showing a hint for a missing member. 209 | missing-member-max-choices=1 210 | 211 | 212 | [LOGGING] 213 | 214 | # Logging modules to check that the string format arguments are in logging 215 | # function parameter format 216 | logging-modules=logging 217 | 218 | 219 | [BASIC] 220 | 221 | # Naming hint for argument names 222 | # argument-name-hint=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$ 223 | argument-name-hint=(([a-z][a-z0-9_]*)|(_[a-z0-9_]*))$ 224 | 225 | # Regular expression matching correct argument names 226 | # argument-rgx=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$ 227 | argument-rgx=(([a-z][a-z0-9_]*)|(_[a-z0-9_]*))$ 228 | 229 | # Naming hint for attribute names 230 | # attr-name-hint=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$ 231 | attr-name-hint=(([a-z][a-z0-9_]*)|(_[a-z0-9_]*))$ 232 | 233 | # Regular expression matching correct attribute names 234 | # attr-rgx=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$ 235 | attr-rgx=(([a-z][a-z0-9_]*)|(_[a-z0-9_]*))$ 236 | 237 | # Bad variable names which should always be refused, separated by a comma 238 | bad-names=foo,bar,baz,toto,tutu,tata 239 | 240 | # Naming hint for class attribute names 241 | # class-attribute-name-hint=([A-Za-z_][A-Za-z0-9_]{2,30}|(__.*__))$ 242 | class-attribute-name-hint=([A-Za-z_][A-Za-z0-9_]*|(__.*__))$ 243 | 244 | # Regular expression matching correct class attribute names 245 | # class-attribute-rgx=([A-Za-z_][A-Za-z0-9_]{2,30}|(__.*__))$ 246 | class-attribute-rgx=([A-Za-z_][A-Za-z0-9_]*|(__.*__))$ 247 | 248 | # Naming hint for class names 249 | class-name-hint=[A-Z_][a-zA-Z0-9]+$ 250 | 251 | # Regular expression matching correct class names 252 | class-rgx=[A-Z_][a-zA-Z0-9]+$ 253 | 254 | # Naming hint for constant names 255 | # const-name-hint=(([A-Z_][A-Z0-9_]*)|(__.*__))$ 256 | const-name-hint=(([a-zA-Z_][a-zA-Z0-9_]*)|(__.*__))$ 257 | 258 | # Regular expression matching correct constant names 259 | # const-rgx=(([A-Z_][A-Z0-9_]*)|(__.*__))$ 260 | const-rgx=(([a-zA-Z_][a-zA-Z0-9_]*)|(__.*__))$ 261 | 262 | # Minimum line length for functions/classes that require docstrings, shorter 263 | # ones are exempt. 264 | docstring-min-length=-1 265 | 266 | # Naming hint for function names 267 | function-name-hint=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$ 268 | # function-name-hint=[A-Z_][a-zA-Z0-9_]*$ 269 | 270 | # Regular expression matching correct function names 271 | function-rgx=(([a-z][a-z0-9_]{2,40})|(_[a-z0-9_]*))$ 272 | # function-rgx=[A-Z_][a-zA-Z0-9_]*$ 273 | 274 | # Good variable names which should always be accepted, separated by a comma 275 | good-names=i,j,k,ex,Run,_ 276 | 277 | # Include a hint for the correct naming format with invalid-name 278 | include-naming-hint=yes 279 | 280 | # Naming hint for inline iteration names 281 | inlinevar-name-hint=[A-Za-z_][A-Za-z0-9_]*$ 282 | 283 | # Regular expression matching correct inline iteration names 284 | inlinevar-rgx=[A-Za-z_][A-Za-z0-9_]*$ 285 | 286 | # Naming hint for method names 287 | method-name-hint=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$ 288 | #method-name-hint=(test|[A-Z_])[a-zA-Z0-9_]*$ 289 | method-naming-style=snake_case 290 | 291 | # Regular expression matching correct method names 292 | method-rgx=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$ 293 | #method-rgx=(test|[A-Z_])[a-zA-Z0-9_]*$ 294 | 295 | # Naming hint for module names 296 | module-name-hint=(([a-z_][a-z0-9_]*)|([A-Z][a-zA-Z0-9]+))$ 297 | 298 | # Regular expression matching correct module names 299 | module-rgx=(([a-z_][a-z0-9_]*)|([A-Z][a-zA-Z0-9]+))$ 300 | 301 | # Colon-delimited sets of names that determine each other's naming style when 302 | # the name regexes allow several styles. 303 | name-group= 304 | 305 | # Regular expression which should only match function or class names that do 306 | # not require a docstring. 307 | no-docstring-rgx=^_ 308 | 309 | # List of decorators that produce properties, such as abc.abstractproperty. Add 310 | # to this list to register other decorators that produce valid properties. 311 | property-classes=abc.abstractproperty 312 | 313 | # Naming hint for variable names 314 | # variable-name-hint=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$ 315 | variable-name-hint=(([a-z][a-z0-9_]*)|(_[a-z0-9_]*))$ 316 | 317 | # Regular expression matching correct variable names 318 | # variable-rgx=(([a-z][a-z0-9_]{2,30})|(_[a-z0-9_]*))$ 319 | variable-rgx=(([a-z][a-z0-9_]*)|(_[a-z0-9_]*))$ 320 | 321 | 322 | [MISCELLANEOUS] 323 | 324 | # List of note tags to take in consideration, separated by a comma. 325 | notes=FIXME,XXX,TODO 326 | 327 | 328 | [FORMAT] 329 | 330 | # Expected format of line ending, e.g. empty (any line ending), LF or CRLF. 331 | expected-line-ending-format= 332 | 333 | # Regexp for a line that is allowed to be longer than the limit. 334 | ignore-long-lines=^\s*(# )??$ 335 | 336 | # Number of spaces of indent required inside a hanging or continued line. 337 | indent-after-paren=4 338 | 339 | # String used as indentation unit. This is usually " " (4 spaces) or "\t" (1 340 | # tab). 341 | # indent-string=' ' 342 | indent-string=' ' 343 | 344 | # Maximum number of characters on a single line. 345 | # max-line-length=100 346 | max-line-length=80 347 | 348 | # Maximum number of lines in a module 349 | max-module-lines=1000 350 | 351 | # List of optional constructs for which whitespace checking is disabled. `dict- 352 | # separator` is used to allow tabulation in dicts, etc.: {1 : 1,\n222: 2}. 353 | # `trailing-comma` allows a space between comma and closing bracket: (a, ). 354 | # `empty-line` allows space-only lines. 355 | no-space-check=trailing-comma,dict-separator 356 | 357 | # Allow the body of a class to be on the same line as the declaration if body 358 | # contains single statement. 359 | single-line-class-stmt=no 360 | 361 | # Allow the body of an if to be on the same line as the test if there is no 362 | # else. 363 | single-line-if-stmt=no 364 | 365 | 366 | [SPELLING] 367 | 368 | # Spelling dictionary name. Available dictionaries: en_US (myspell). 369 | spelling-dict= 370 | 371 | # List of comma separated words that should not be checked. 372 | spelling-ignore-words= 373 | 374 | # A path to a file that contains private dictionary; one word per line. 375 | spelling-private-dict-file= 376 | 377 | # Tells whether to store unknown words to indicated private dictionary in 378 | # --spelling-private-dict-file option instead of raising a message. 379 | spelling-store-unknown-words=no 380 | 381 | 382 | [SIMILARITIES] 383 | 384 | # Ignore comments when computing similarities. 385 | ignore-comments=yes 386 | 387 | # Ignore docstrings when computing similarities. 388 | ignore-docstrings=yes 389 | 390 | # Ignore imports when computing similarities. 391 | ignore-imports=no 392 | 393 | # Minimum lines number of a similarity. 394 | min-similarity-lines=4 395 | 396 | 397 | [DESIGN] 398 | 399 | # Maximum number of arguments for function / method 400 | # max-args=5 401 | max-args=10 402 | 403 | # Maximum number of attributes for a class (see R0902). 404 | max-attributes=7 405 | 406 | # Maximum number of boolean expressions in a if statement 407 | max-bool-expr=5 408 | 409 | # Maximum number of branch for function / method body 410 | max-branches=12 411 | 412 | # Maximum number of locals for function / method body 413 | max-locals=15 414 | 415 | # Maximum number of parents for a class (see R0901). 416 | max-parents=7 417 | 418 | # Maximum number of public methods for a class (see R0904). 419 | max-public-methods=20 420 | 421 | # Maximum number of return / yield for function / method body 422 | max-returns=6 423 | 424 | # Maximum number of statements in function / method body 425 | max-statements=50 426 | 427 | # Minimum number of public methods for a class (see R0903). 428 | min-public-methods=2 429 | 430 | 431 | [CLASSES] 432 | 433 | # List of method names used to declare (i.e. assign) instance attributes. 434 | defining-attr-methods=__init__,__new__,setUp 435 | 436 | # List of member names, which should be excluded from the protected access 437 | # warning. 438 | exclude-protected=_asdict,_fields,_replace,_source,_make 439 | 440 | # List of valid names for the first argument in a class method. 441 | valid-classmethod-first-arg=cls 442 | 443 | # List of valid names for the first argument in a metaclass class method. 444 | valid-metaclass-classmethod-first-arg=mcs 445 | 446 | 447 | [IMPORTS] 448 | 449 | # Allow wildcard imports from modules that define __all__. 450 | allow-wildcard-with-all=no 451 | 452 | # Analyse import fallback blocks. This can be used to support both Python 2 and 453 | # 3 compatible code, which means that the block might have code that exists 454 | # only in one or another interpreter, leading to false positives when analysed. 455 | analyse-fallback-blocks=no 456 | 457 | # Deprecated modules which should not be used, separated by a comma 458 | deprecated-modules=optparse,tkinter.tix 459 | 460 | # Create a graph of external dependencies in the given file (report RP0402 must 461 | # not be disabled) 462 | ext-import-graph= 463 | 464 | # Create a graph of every (i.e. internal and external) dependencies in the 465 | # given file (report RP0402 must not be disabled) 466 | import-graph= 467 | 468 | # Create a graph of internal dependencies in the given file (report RP0402 must 469 | # not be disabled) 470 | int-import-graph= 471 | 472 | # Force import order to recognize a module as part of the standard 473 | # compatibility libraries. 474 | known-standard-library= 475 | 476 | # Force import order to recognize a module as part of a third party library. 477 | known-third-party=enchant 478 | 479 | 480 | [EXCEPTIONS] 481 | 482 | # Exceptions that will emit a warning when being caught. Defaults to 483 | # "Exception" 484 | overgeneral-exceptions=Exception 485 | -------------------------------------------------------------------------------- /.style.yapf: -------------------------------------------------------------------------------- 1 | # For documentation of these options see 2 | # https://github.com/google/yapf/blob/HEAD/README.rst 3 | # based on Fuchsia YAPF config: 4 | # https://fuchsia.googlesource.com/fuchsia/+/refs/heads/main/.style.yapf 5 | [style] 6 | based_on_style = google 7 | indent_dictionary_value = true 8 | split_arguments_when_comma_terminated = true 9 | split_before_first_argument = true 10 | indent_width = 2 11 | column_limit = 80 12 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "editor.formatOnSave": true, 3 | "python.analysis.typeCheckingMode": "strict", 4 | "python.testing.pytestArgs": [ 5 | "picatrix" 6 | ], 7 | "python.testing.unittestEnabled": false, 8 | "python.testing.pytestEnabled": true, 9 | "python.testing.unittestArgs": [ 10 | "-v", 11 | "-s", 12 | "./picatrix", 13 | "-p", 14 | "*_test.py" 15 | ], 16 | "python.formatting.provider": "yapf", 17 | "python.languageServer": "Pylance", 18 | "python.linting.enabled": true, 19 | "python.linting.pylintEnabled": true, 20 | "[python]": { 21 | "editor.codeActionsOnSave": { 22 | "source.organizeImports": true 23 | } 24 | }, 25 | "files.insertFinalNewline": true, 26 | } 27 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright [yyyy] [name of copyright owner] 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. 203 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Picatrix 2 | 3 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/picatrix/blob/main/notebooks/Quick_Primer_on_Colab_Jupyter.ipynb) 4 | [![Open In Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/google/picatrix.git/main?filepath=notebooks%2F) 5 | [![Version](https://img.shields.io/pypi/v/picatrix.svg)](https://pypi.python.org/pypi/picatrix) 6 | ![GitHub e2e Test Status](https://img.shields.io/github/workflow/status/google/picatrix/picatrix-end-to-end) 7 | 8 | Picatrix is a framework that is meant to be used within a [Colab](https://colab.research.google.com) or 9 | [Jupyter](https://jupyter.org/) notebooks. The framework is designed around 10 | providing a security analyst with the libraries to develop helper functions 11 | that will be exposed as magics and regular python functions in notebooks. 12 | 13 | This makes it easier to share an environment with other analysts, exposing 14 | common functions that are used in notebooks to everyone. In addition to that 15 | the functions themselves are designed to make it easier to work with various 16 | APIs and backends in a notebook environment. The functions mostly involve 17 | returning data back as a pandas DataFrame for further processing or to work 18 | with pandas (manipulate pandas, change values, enrich data, upload data frames 19 | to other services, etC). 20 | 21 | ## Howto Get Started 22 | 23 | Read the [installation instructions](docs/Installation.md) on the best ways 24 | to install picatrix. 25 | 26 | After installing, connect to the Jupyter notebook in your web browser (should open 27 | up automatically). Inside the notebook you need to import the picatrix library 28 | and initialize it: 29 | 30 | ``` 31 | from picatrix import notebook_init 32 | notebook_init.init() 33 | ``` 34 | 35 | (if you are using the docker container you don't need to import these libraries, 36 | that is done for you automatically). 37 | 38 | And that's it, then all the magics/helper functions are now ready and accessible 39 | to your notebook. To get a list of the available helpers, use: 40 | 41 | ``` 42 | %picatrixmagics 43 | ``` 44 | 45 | Or 46 | 47 | ``` 48 | picatrixmagics_func() 49 | ``` 50 | 51 | Each magic has a `--help` parameter or the functions with `_func?`. Eg. 52 | 53 | ``` 54 | timesketch_set_active_sketch_func? 55 | ``` 56 | 57 | ## Examples 58 | 59 | To get all sketches, you can use the following magic 60 | 61 | ``` 62 | %timesketch_get_sketches 63 | ``` 64 | 65 | For most of the magics you need to set an active sketch 66 | 67 | ``` 68 | %timesketch_set_active_sketch 1 69 | ``` 70 | 71 | To query the sketch, the following magic will execute a search and return the results as a search object, 72 | that can be easily converted into a pandas dataframe: 73 | 74 | ``` 75 | search_obj = %timesketch_query 'message:Another' 76 | search_obj.table 77 | ``` 78 | 79 | Further documentation on the search object can be [found 80 | here](https://timesketch.org/developers/api-client/#search-query) 81 | 82 | 83 | To add a manual event with a function use: 84 | 85 | ``` 86 | timesketch_add_manual_event_func('Eventdescriptiontext', attributes=attributesdict) 87 | ``` 88 | 89 | Which is the same as: 90 | ``` 91 | %timesketch_add_manual_event Eventdescriptiontext --attributes {{attributesdict}} 92 | ``` 93 | 94 | ## Discussions 95 | 96 | Want to discuss the project, have issues, want new features, join the slack 97 | workspace [here](http://join-open-source-dfir-slack.herokuapp.com/), the 98 | channel for picatrix is #picatrix. 99 | -------------------------------------------------------------------------------- /docker/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM python:3.8-slim 2 | 3 | # Create folders and fix permissions. 4 | RUN groupadd --gid 1000 picagroup && \ 5 | useradd picatrix --uid 1000 --gid 1000 -d /home/picatrix -m && \ 6 | mkdir -p /usr/local/src/picadata/ && \ 7 | chmod 777 /usr/local/src/picadata/ 8 | 9 | USER picatrix 10 | WORKDIR /home/picatrix 11 | ENV VIRTUAL_ENV=/home/picatrix/picenv 12 | 13 | RUN python3 -m venv $VIRTUAL_ENV && \ 14 | mkdir -p .ipython/profile_default/startup/ && \ 15 | mkdir -p /home/picatrix/.jupyter 16 | 17 | ENV PATH="$VIRTUAL_ENV/bin:$PATH" 18 | ENV JUPYTER_PORT=8899 19 | 20 | COPY --chown=1000:1000 docker/notebook_init.py /home/picatrix/.ipython/profile_default/startup/notebook_init.py 21 | COPY --chown=1000:1000 . /home/picatrix/code 22 | COPY --chown=1000:1000 docker/jupyter_notebook_config.py /home/picatrix/.jupyter/jupyter_notebook_config.py 23 | 24 | RUN pip install --upgrade pip setuptools wheel && \ 25 | cd /home/picatrix/code && pip install -e .[runtime] && \ 26 | bash prepare-picatrix-runtime.sh 27 | 28 | WORKDIR /usr/local/src/picadata/ 29 | EXPOSE 8899 30 | 31 | # Run jupyter. 32 | ENTRYPOINT ["jupyter", "notebook"] 33 | -------------------------------------------------------------------------------- /docker/README.md: -------------------------------------------------------------------------------- 1 | # Docker for picatrix 2 | 3 | Picatrix has support for Docker. This is a convenient way of getting up and running. 4 | 5 | ### Install Docker 6 | Follow the official instructions [here](https://www.docker.com/community-edition) 7 | 8 | ### Install Docker Compose 9 | Follow the official instructions [here](https://docs.docker.com/compose/install/) 10 | 11 | ### Clone Picatrix 12 | 13 | ```shell 14 | $ git clone https://github.com/google/picatrix.git 15 | cd picatrix/docker 16 | ``` 17 | 18 | ### Build and Start Containers 19 | 20 | ```shell 21 | $ sudo docker-compose -f docker-build.yml build 22 | $ sudo docker-compose -f docker-build.yml up -d 23 | ``` 24 | 25 | To build the container use: 26 | 27 | ```shell 28 | $ sudo docker-compose --env-file config.env build 29 | ``` 30 | 31 | ### Access Picatrix 32 | 33 | To access picatrix you need to start a browser and paste in the following 34 | URL: http://localhost:8899/?token=picatrix 35 | 36 | And you should have a working Jupyter installation with picatrix ready. 37 | 38 | If you want to change the token, edit the file `docker/jupyter_notebook_config.py` 39 | before building the docker image. 40 | 41 | ### Can I get Access to Data? 42 | 43 | The /tmp/ folder is mounted as data inside the notebook container. Therefore if 44 | you need to import data from your local machine, make a copy of it available in 45 | the /tmp/ folder and read it from there (or change the docker config files to 46 | expose other folders). 47 | -------------------------------------------------------------------------------- /docker/docker-build.yml: -------------------------------------------------------------------------------- 1 | version: '3.6' 2 | services: 3 | picatrix: 4 | build: 5 | context: ../ 6 | dockerfile: ./docker/Dockerfile 7 | ports: 8 | - 8899:8899 9 | restart: on-failure 10 | volumes: 11 | - ../:/usr/local/src/picatrix/:ro 12 | - /tmp/:/usr/local/src/picadata/ 13 | -------------------------------------------------------------------------------- /docker/docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: '3.6' 2 | services: 3 | picatrix: 4 | image: us-docker.pkg.dev/osdfir-registry/picatrix/picatrix:latest 5 | ports: 6 | - 8899:8899 7 | restart: on-failure 8 | volumes: 9 | - ../:/usr/local/src/picatrix/:ro 10 | - /tmp/:/usr/local/src/picadata/ 11 | -------------------------------------------------------------------------------- /docker/jupyter_notebook_config.py: -------------------------------------------------------------------------------- 1 | # Configuration file for jupyter-notebook. 2 | 3 | ## Use a regular expression for the Access-Control-Allow-Origin header 4 | # 5 | # Requests from an origin matching the expression will get replies with: 6 | # 7 | # Access-Control-Allow-Origin: origin 8 | # 9 | # where `origin` is the origin of the request. 10 | # 11 | # Ignored if allow_origin is set. 12 | c.NotebookApp.allow_origin_pat = 'https://colab.[a-z]+.google.com' 13 | 14 | ## The IP address the notebook server will listen on. 15 | c.NotebookApp.ip = '*' 16 | 17 | ## The directory to use for notebooks and kernels. 18 | # Uncomment this if you want the notebook to start immediately in this folder. 19 | #c.NotebookApp.notebook_dir = 'data/' 20 | 21 | ## Whether to open in a browser after starting. The specific browser used is 22 | # platform dependent and determined by the python standard library `webbrowser` 23 | # module, unless it is overridden using the --browser (NotebookApp.browser) 24 | # configuration option. 25 | c.NotebookApp.open_browser = False 26 | 27 | ## Hashed password to use for web authentication. 28 | # 29 | # To generate, type in a python/IPython shell: 30 | # 31 | # from notebook.auth import passwd; passwd() 32 | # 33 | # The string should be of the form type:salt:hashed-password. 34 | # If this is enabled the password "picatrix" can be used. 35 | #c.NotebookApp.password = 'argon2:$argon2id$v=19$m=10240,t=10,p=8$Q2sRv8dVZ8WBSmGNcTMuKg$VtFU0bwX81Ou+OaDWQgloA' 36 | # Right now the token is set to picatrix, a plain text password. 37 | c.NotebookApp.token = 'picatrix' 38 | 39 | ## The port the notebook server will listen on. 40 | c.NotebookApp.port = 8899 41 | 42 | ## The number of additional ports to try if the specified port is not available. 43 | c.NotebookApp.port_retries = 0 44 | 45 | ## The base name used when creating untitled notebooks. 46 | c.ContentsManager.untitled_notebook = 'NewPicatrixNotebook' 47 | -------------------------------------------------------------------------------- /docker/notebook_init.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """This is an import file that runs on every startup of the Jupyter runtime.""" 15 | 16 | from picatrix import notebook_init 17 | 18 | notebook_init.init() 19 | -------------------------------------------------------------------------------- /docs/AddingAMagic.md: -------------------------------------------------------------------------------- 1 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/picatrix/blob/main/notebooks/adding_magic.ipynb) 2 | 3 | # Adding a Magic 4 | 5 | There are two ways of adding a new magic or a helper function into 6 | the picatrix library. 7 | 8 | 1. Temporary registration in a notebook. 9 | 2. Permanent one by checking files into the codebase. 10 | 11 | Let's explore both. 12 | 13 | There is also a second way to explore this document or these instructions, 14 | through this document or through the [interactive notebook](https://colab.research.google.com/github/google/picatrix/blob/main/notebooks/adding_magic.ipynb) 15 | 16 | ## Temporary Registration in a Notebook 17 | 18 | Whether there is a need to register a helper function that will be 19 | very specific to a single notebook and not helpful for the greater 20 | community or in order to test a function while developin it, having 21 | the ability to temporarily register a helper function inside a notebook 22 | can be very handy. In order to do that you'll need two things: 23 | 24 | 1. Import the framework library 25 | 2. Add a decorator to a function. 26 | 27 | The framework library needs to be loaded in order to register a magic, 28 | that can be simply done: 29 | 30 | ``` 31 | from picatrix.lib import framework 32 | ``` 33 | 34 | Then a magic can be registered. A magic or a helper function is a regular 35 | Python function that accepts at least one parameter, `data`. Other parameters 36 | are optional and will be translated into regular parameters to the magic 37 | and/or function. The function needs to have a proper docstring and typing 38 | in order to correctly register. Then all that is needed in order to pass 39 | it in is to put the decorator `@framework.picatrix_magic` on the function. 40 | 41 | An example would be: 42 | 43 | ``` 44 | @framework.picatrix_magic 45 | def my_silly_magic(data: Text, magnitude: Optional[int] = 100) -> Text: 46 | """Returns a string to demonstrate how to construct a magic. 47 | 48 | Args: 49 | data (str): This is a string that will be printed back. 50 | magnitude (int): A number that will be displayed in the string. 51 | 52 | Returns: 53 | A string that basically combines the two options. 54 | """ 55 | return f'This magical magic produced {magnitude} magics of {data.strip()}' 56 | ``` 57 | 58 | As you can see in the above description there is a function declaration (def my...) 59 | that defines the name of the magic. The magic will now be registered as 60 | `%my_silly_magic`, `%%my_silly_magic` and `my_silly_magic_func()`. In both cell 61 | and line mode the text that is added will be passed in as the data attribute. 62 | That is defined as a `str` (both by typing and in args section). That means the 63 | input value will be interpreted as a string. You can pass in the optional 64 | parameter `magnitude`, but by default it will be set to 100. So this could 65 | be run as: 66 | 67 | ``` 68 | %my_silly_magic what today is magical 69 | ``` 70 | 71 | And that would return back: 72 | ``` 73 | This magical magic produced 100 magics of what today is magical 74 | ``` 75 | 76 | You can also run the function as: 77 | 78 | ``` 79 | %%my_silly_magic foobar 80 | magics of what today is magical 81 | this is surely a sorrow of a magic 82 | ``` 83 | 84 | And that would store the results of the magic into the variable `foobar`, so that 85 | it would be: 86 | 87 | ``` 88 | foobar == 'This magical magic produced 100 magics of magics of what today is magical\nthis is surely a sorrow of a magic' 89 | ``` 90 | 91 | To pass in the parameters you can do: 92 | 93 | ``` 94 | my_silly_magic('foo', magnitude=1) 95 | ``` 96 | 97 | That would produce: 98 | 99 | ``` 100 | This magical magic produced 1 magics of foo 101 | ``` 102 | 103 | Or you can do it this way: 104 | 105 | ``` 106 | %my_silly_magic --bindto bar --magnitude 203 foobar 107 | ``` 108 | 109 | That would store the results into a variable called `bar`, so that: 110 | 111 | ``` 112 | bar == 'This magical magic produced 203 magics of foobar 113 | ``` 114 | 115 | This is all that is required to register the magic temporarily. 116 | 117 | ## Checking the Magic Into the Library 118 | 119 | If you want the magic to be more permanent the magic definition needs 120 | to be checked into the library. 121 | 122 | 1. Create an appropriately named file inside `picatrix/magics/FILE.py` 123 | 2. Add an import statement for the file into `picatrix/magics/__init__.py` 124 | 3. Add unit tests for the magics into `picatrix/magics/FILE_test.py` 125 | 4. If needed, add e2e tests to the magic. 126 | 5. If needed, add new dependencies into `picatrix/dependencies.py`. 127 | -------------------------------------------------------------------------------- /docs/Contributing.md: -------------------------------------------------------------------------------- 1 | # How to Contribute 2 | 3 | We'd love to accept your patches and contributions to this project. There are 4 | just a few small guidelines you need to follow. 5 | 6 | ## Contributor License Agreement 7 | 8 | Contributions to this project must be accompanied by a Contributor License 9 | Agreement. You (or your employer) retain the copyright to your contribution; 10 | this simply gives us permission to use and redistribute your contributions as 11 | part of the project. Head over to to see 12 | your current agreements on file or to sign a new one. 13 | 14 | You generally only need to submit a CLA once, so if you've already submitted one 15 | (even if it was for a different project), you probably don't need to do it 16 | again. 17 | 18 | ## Code reviews 19 | 20 | To keep the code base maintainable and readable all code is developed using a similar coding style. See the [style guide](https://google.github.io/styleguide/pyguide.html). This makes the code easier to maintain and understand. 21 | 22 | The purpose of the code review is to ensure that: 23 | 24 | * at least two eyes looked over the code in hopes of finding potential bugs or errors (before they become bugs and errors). This also improves the overall code quality. 25 | * make sure the code adheres to the style guide. 26 | * review design decisions and if needed assist with making the code more optimal or error tolerant. 27 | 28 | We use GitHub pull requests for this purpose. Consult 29 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more 30 | information on using pull requests. 31 | 32 | **The short version:** 33 | 34 | *don't be intimidated.* 35 | 36 | **The longer version:** 37 | 38 | One language is not the same as another, you might be fluent in C or Perl, but that does not mean the same fluency for Python. You might have just started programming while others have been doing this for years. Our challenge is having a code base that is accessible and sufficiently uniform to most of you. 39 | 40 | Also don't be intimidated by rewrites/refactors, which can make it feel like the code base is changing under your feet. We have to make sure the code base is maintainable and this involves regular reshaping and cleaning up to get new features in. 41 | 42 | We continuously try to improve the code base, including making things and easier and quicker to write. This sometimes means that the way you just learned might already superseded by another. We try to keep the documentation up to date but this will sometimes be after you ran into an issue. 43 | 44 | First time contributors may come across the fact that the code review process can take quite a long time, with lots of back and forth comments. You may think that you are wasting the core developers time, but rest assured you are not. We look at this as an investment of building up good solid code contributors. We want to make sure our contributors understand the code and the style guide and will make suggestions to the contributor to fix what we think needs improving. Despite spending potentially more time to begin with to get code submitted into the project we believe this investment in code review will result in better code submissions and increased proficiency of the contributor. 45 | 46 | Therefore we would like to ask people to hang on, to get through the code review process and try to learn something while going through it. Rest assured, it will get easier next time and even easier the time after that, and before you know it you can contribute code to the project with little to no comments. 47 | 48 | And if things are unclear, don't hesitate to ask. You can also join the slack channel [here](http://join-open-source-dfir-slack.herokuapp.com/), the channel for picatrix is #picatrix. 49 | 50 | ## Community Guidelines 51 | 52 | This project follows [Google's Open Source Community 53 | Guidelines](https://opensource.google/conduct/). 54 | -------------------------------------------------------------------------------- /docs/DevelopmentEnvironment.md: -------------------------------------------------------------------------------- 1 | # Setting up Your Development Environment 2 | 3 | This guide is to help you set up your development environment for picatrix. 4 | 5 | The general flow is: 6 | 7 | 1. Fork the project into your personal account 8 | 2. Grab a copy of the personal fork to your machine (git clone) 9 | 3. Create a feature branch on the personal fork 10 | 4. Work on your code, commit the code to the feature branch 11 | 5. Push the feature branch to your personal fork 12 | 6. Create a pull request (PR) on the main project. 13 | 14 | ## Setting up Git 15 | 16 | Start by visiting the [project](https://github.com/google/picatrix) and click the `Fork` 17 | button in the upper right corner of the project to fork the project into your personal 18 | account. 19 | 20 | Once the project has been forked you need to clone the personal fork and add the upstream project. 21 | 22 | ```shell 23 | $ git clone https://github.com//picatrix.git 24 | $ cd picatrix 25 | $ git remote add upstream https://github.com/google/picatrix.git 26 | ``` 27 | 28 | You also need to make sure that git is correctly configured 29 | 30 | ```shell 31 | $ git config --global user.name "Full Name" 32 | $ git config --global user.email name@example.com 33 | $ git config --global push.default matching 34 | ``` 35 | 36 | Next is to make sure that `.netrc` is configured. For more information see: https://gist.github.com/technoweenie/1072829 37 | 38 | ## Create a Feature Branch 39 | 40 | Every code change is stored in a separate feature branch: 41 | 42 | ### Sync Main Branch 43 | 44 | Make sure the master branch of your fork is synced with upstream: 45 | 46 | ```shell 47 | $ git checkout main 48 | $ git fetch upstream 49 | $ git rebase upstream/main 50 | $ git push 51 | ``` 52 | 53 | or in a single command: 54 | ```shell 55 | $ git checkout main && git fetch upstream && git rebase upstream/main && git push 56 | ``` 57 | 58 | ### Create a feature branch: 59 | 60 | ```shell 61 | $ git checkout -b 62 | ``` 63 | 64 | eg: 65 | 66 | ```shell 67 | $ git checkout -b my_awesome_feature 68 | ``` 69 | 70 | Make the necessary code changes, commit them to your feature branch and upload them to your fork: 71 | 72 | ```shell 73 | $ git push origin 74 | ``` 75 | 76 | ## Make Changes 77 | 78 | Make your changes to the code and make sure to push them to the feature branch of your fork. 79 | 80 | Commit all code changes to your feature branch and upload them: 81 | 82 | ```shell 83 | $ git push origin 84 | ``` 85 | 86 | ## Start a code review 87 | 88 | Once you think your changes are ready, you start the review process. There are few ways of doing that, 89 | one is to use the [Github web UI](https://docs.github.com/en/free-pro-team@latest/github/collaborating-with-issues-and-pull-requests/creating-a-pull-request). 90 | Visit the [project](https://github.com/google/picatrix/pulls) and click the `create a pull request` button. 91 | 92 | The other option is to use the Github CLI tool. 93 | 94 | Select a good descriptive PR name and write a description describing the change. After all is completed, select a code reviewer 95 | from the list of available reviewers (if you don't have the option to select a reviewer, one will be 96 | assigned). 97 | 98 | ## Status checks 99 | Once your pull request has been created, automated checkers will run to on your changes to check 100 | for mistakes, or code that doesn't match the style guide. Review the output from the tools, and 101 | make sure your pull request passes. 102 | 103 | Your pull request cannot be merged until the checkers report everything is OK. 104 | 105 | ## Making changes after reviews 106 | 107 | Your pull request will be assigned to a project maintainer either by you or the reviewer. 108 | The maintainer will review your code, and may push some changes to your branch and/or request 109 | that you make changes. If they do make changes, make sure your local copy of the branch 110 | (git pull) before making further changes. 111 | 112 | After they've reviewed your code, make any necessary changes, and push them to your 113 | feature branch. After that, reply to the comments made by the reviewer, then request 114 | a new review from the same reviewer. In the upper right corner in the `Reviewers` 115 | category a button should be there to re-request a review from a reviewer (ATM it is a circle, 116 | or arrows that make up a circle). 117 | 118 | ## Merging 119 | Once the pull request assignee is happy with your proposed changes, and all the status 120 | checks have passed, the assignee will merge your pull request into the project. 121 | Once that's completed, you can delete your feature branch. 122 | -------------------------------------------------------------------------------- /docs/Installation.md: -------------------------------------------------------------------------------- 1 | # Installation 2 | 3 | Picatrix can be installed via two ways: 4 | 5 | 1. docker 6 | 2. pip inside a virtualenv 7 | 8 | Let's explore both methods: 9 | 10 | ## Docker 11 | 12 | The easiest way to install picatrix is via docker. To use it, first 13 | [install docker](https://docs.docker.com/engine/install/). 14 | 15 | If you did not install the docker desktop app you may also need to install 16 | `docker-compose`, please follow the instructions 17 | [here](https://docs.docker.com/compose/install/) (the version that is often 18 | included in your source repo might be too old to properly setup the container). 19 | 20 | After installing docker the next step is to download the source code: 21 | 22 | ```shell 23 | $ git clone https://github.com/google/picatrix.git 24 | ``` 25 | 26 | *(if you don't have the git client installed you can also download 27 | the [source code using this link](https://github.com/google/picatrix/archive/main.zip))* 28 | 29 | ### Install via Bash Script 30 | 31 | If you are running on a Linux/Mac you may try to first run the installation 32 | script: 33 | 34 | ```shell 35 | $ sh install.sh 36 | ``` 37 | 38 | This script will create a folder `${HOME}/picadata` that will be used as a 39 | persistent storage for the picatrix container and fetch and start the container. 40 | 41 | It will also create a new file, called `docker/docker-tmp.yml` that defines 42 | the new location of the mapped data. 43 | 44 | If for some reasons the you are not able to create new notebooks, or see 45 | the example notebooks (this may happen on a Mac OS X for instance). 46 | Please edit the `docker/docker-tmp.yml` file and change the `~/picadata` 47 | line to a full path, eg. `/home/foobar/picadata` and run: 48 | 49 | ```shell 50 | $ sudo docker container stop docker_picatrix_1 51 | $ sudo docker-compose -f docker/docker-tmp.yml up -d 52 | ``` 53 | 54 | ### Default Docker Script 55 | 56 | If you are not running on a Linux/Mac machine or want to customize the 57 | installation yourself you can run: 58 | 59 | ```shell 60 | $ cd picatrix/docker 61 | $ sudo docker-compose up -d 62 | ``` 63 | 64 | Please note that by default the /tmp folder on your host will be mapped into 65 | a `data` folder on the docker container. If you want to change that and point 66 | to another folder on your system (**highly encouraged**), edit the file 67 | `docker-latest.yml` and change the path `/tmp` to a folder of your choosing 68 | (just remember that the folder needs to be writable by uid=1000 and/or 69 | gid=1000 if you are running a Linux based host). 70 | 71 | For instance if you are running this on a Windows system, then you will 72 | need to change the `/tmp/` to something like `C:\My Folder\Where I store data`. 73 | Also when running on Windows, there is no `sudo` in front of the commands. 74 | 75 | The docker compose command will download the latest build and deploy the 76 | picatrix docker container. 77 | 78 | 79 | ### Access the Container 80 | 81 | To be able to connect to picatrix, open the following URL in a browser 82 | window: 83 | [http://localhost:8899/?token=picatrix](http://localhost:8899/?token=picatrix)` 84 | 85 | You should have a fully working copy of Jupyter running, with an 86 | active picatrix library. 87 | 88 | You can also just visit [http://localhost:8899](http://localhost:8899) and 89 | type in `picatrix` as the password. 90 | 91 | ### Upgrade Container 92 | 93 | Depending on whether the docker container got deployed using the default 94 | docker container configuration file or the one that was generated by the 95 | `install.sh` script run the following commands: 96 | 97 | #### Install Script 98 | 99 | ```shell 100 | $ cd picatrix/docker 101 | $ sudo docker-compose -f docker-tmp.yml pull 102 | $ sudo docker-compose -f docker-tmp.yml up -d 103 | ``` 104 | 105 | #### Default Config 106 | 107 | ```shell 108 | $ cd picatrix/docker 109 | $ sudo docker-compose pull 110 | $ sudo docker-compose up -d 111 | ``` 112 | 113 | #### Manually Update 114 | 115 | You can also manually pull the new image using: 116 | 117 | ```shell 118 | $ sudo docker pull us-docker.pkg.dev/osdfir-registry/picatrix/picatrix:latest 119 | ``` 120 | 121 | #### Docker Desktop 122 | 123 | If you are using Docker desktop you can find the docker image, click 124 | on the three dots and select pull. 125 | 126 | After manually updating the image the container needs to be recreated (using 127 | the docker compose up command used earlier). 128 | 129 | ## Virtualenv 130 | 131 | To install picatrix using virtualenv you can use the following commands: 132 | 133 | ```shell 134 | $ python3 -m venv picatrix_env 135 | $ source picatrix_env/bin/activate 136 | $ pip install picatrix 137 | ``` 138 | 139 | And then to start picatrix you can either run: 140 | 141 | ```shell 142 | $ jupyter notebook 143 | ``` 144 | 145 | If you want to connect from a colab frontend, then you'll need to run these 146 | two commands as well: 147 | 148 | ```shell 149 | $ pip install jupyter_http_over_ws 150 | $ jupyter serverextension enable --py jupyter_http_over_ws 151 | ``` 152 | 153 | And then to run the notebook: 154 | 155 | ```shell 156 | $ jupyter notebook \ 157 | --NotebookApp.allow_origin='https://colab.research.google.com' \ 158 | --port=8888 \ 159 | --NotebookApp.port_retries=0 160 | ``` 161 | 162 | After the notebook is started you need to load up picatrix using the code cell: 163 | 164 | ```python 165 | from picatrix import notebook_init 166 | notebook_init.init() 167 | ``` 168 | 169 | *This also works if you want to connect to the hosted colab runtime. There you 170 | can simply add a cell with `!pip install --upgrade picatrix` and you should 171 | be able to start using picatrix.* 172 | 173 | ## Confirm Installation 174 | 175 | You can confirm the successful installation by creating a new notebook and type in: 176 | ``` 177 | %picatrixmagics 178 | ``` 179 | 180 | Picatrix library is already imported and initialized. 181 | 182 | ## Connect To Colab 183 | 184 | ### Using Docker Container 185 | 186 | In order to connect to the docker container from colab, select the arrow 187 | next to the `Connect` button, select `Connect to local runtime` and type 188 | in the URL `http://localhost:8899/?token=picatrix` into the `Backend URL` 189 | field and hit `CONNECT`. 190 | 191 | ### Using Virtualenv 192 | 193 | Look at the window where the jupyter notebook is being run from. It should 194 | have lines similar to this: 195 | 196 | ``` 197 | To access the notebook, open this file in a browser: 198 | ... 199 | Or copy and paste one of these URLs: 200 | http://....:8899/?token=... 201 | or http://127.0.0.1:8899/?token=... 202 | ``` 203 | 204 | Copy the URL that is provided there, replace 127.0.0.1 with localhost and enter 205 | the URL into the `Backend URL` as described in the Docker instructions. 206 | -------------------------------------------------------------------------------- /docs/code-of-conduct.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to making participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, gender identity and expression, level of 9 | experience, education, socio-economic status, nationality, personal appearance, 10 | race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or reject 41 | comments, commits, code, wiki edits, issues, and other contributions that are 42 | not aligned to this Code of Conduct, or to ban temporarily or permanently any 43 | contributor for other behaviors that they deem inappropriate, threatening, 44 | offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies both within project spaces and in public spaces 49 | when an individual is representing the project or its community. Examples of 50 | representing a project or community include using an official project e-mail 51 | address, posting via an official social media account, or acting as an appointed 52 | representative at an online or offline event. Representation of a project may be 53 | further defined and clarified by project maintainers. 54 | 55 | This Code of Conduct also applies outside the project spaces when the Project 56 | Steward has a reasonable belief that an individual's behavior may have a 57 | negative impact on the project or its community. 58 | 59 | ## Conflict Resolution 60 | 61 | We do not believe that all conflict is bad; healthy debate and disagreement 62 | often yield positive results. However, it is never okay to be disrespectful or 63 | to engage in behavior that violates the project’s code of conduct. 64 | 65 | If you see someone violating the code of conduct, you are encouraged to address 66 | the behavior directly with those involved. Many issues can be resolved quickly 67 | and easily, and this gives people more control over the outcome of their 68 | dispute. If you are unable to resolve the matter for any reason, or if the 69 | behavior is threatening or harassing, report it. We are dedicated to providing 70 | an environment where participants feel welcome and safe. 71 | 72 | Reports should be directed to *[PROJECT STEWARD NAME(s) AND EMAIL(s)]*, the 73 | Project Steward(s) for *[PROJECT NAME]*. It is the Project Steward’s duty to 74 | receive and address reported violations of the code of conduct. They will then 75 | work with a committee consisting of representatives from the Open Source 76 | Programs Office and the Google Open Source Strategy team. If for any reason you 77 | are uncomfortable reaching out to the Project Steward, please email 78 | opensource@google.com. 79 | 80 | We will investigate every complaint, but you may not receive a direct response. 81 | We will use our discretion in determining when and how to follow up on reported 82 | incidents, which may range from not taking action to permanent expulsion from 83 | the project and project-sponsored spaces. We will notify the accused of the 84 | report and provide them an opportunity to discuss it before any action is taken. 85 | The identity of the reporter will be omitted from the details of the report 86 | supplied to the accused. In potentially harmful situations, such as ongoing 87 | harassment or threats to anyone's safety, we may take action without notice. 88 | 89 | ## Attribution 90 | 91 | This Code of Conduct is adapted from the Contributor Covenant, version 1.4, 92 | available at 93 | https://www.contributor-covenant.org/version/1/4/code-of-conduct.html 94 | -------------------------------------------------------------------------------- /end_to_end_tests/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google Inc. All rights reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """End to end test module.""" 15 | 16 | # Register all tests by importing them. 17 | from . import basic_test, timesketch 18 | -------------------------------------------------------------------------------- /end_to_end_tests/basic_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google Inc. All rights reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """End to end tests of common picatrix magics.""" 15 | 16 | from IPython.terminal.interactiveshell import TerminalInteractiveShell 17 | 18 | from . import interface, manager 19 | 20 | MAGIC_DEFINITION = ( 21 | 'from typing import Optional\n' 22 | 'from typing import Text\n' 23 | '\n' 24 | 'from picatrix.lib import framework\n' 25 | '\n' 26 | '@framework.picatrix_magic\n' 27 | 'def my_silly_magic(data: Text, magnitude: Optional[int] = 100) -> Text:\n' 28 | ' """Return a silly string with no meaningful value.\n' 29 | '\n' 30 | ' Args:\n' 31 | ' data (str): This is a string that will be printed back.\n' 32 | ' magnitude (int): A number that will be displayed in the string.\n' 33 | '\n' 34 | ' Returns:\n' 35 | ' A string that basically combines the two options.\n' 36 | ' """\n' 37 | ' return f"This magical magic produced {magnitude} magics of ' 38 | '{data.strip()}"\n') 39 | 40 | 41 | class BasicTest(interface.BaseEndToEndTest): 42 | """End to end tests for query functionality.""" 43 | 44 | NAME = 'basic_test' 45 | 46 | def test_picatrixmagics(self, ip: TerminalInteractiveShell): 47 | """Test the picatrixmagics.""" 48 | magics = ip.run_line_magic(magic_name='picatrixmagics', line='') 49 | 50 | self.assertions.assertFalse(magics.empty) 51 | self.assertions.assertTrue(magics.shape[0] > 10) 52 | 53 | def test_magic_registration(self, ip: TerminalInteractiveShell): 54 | """Test registering a magic.""" 55 | res = ip.run_cell(raw_cell=MAGIC_DEFINITION) 56 | self.assertions.assertTrue(res.success) 57 | 58 | magics = ip.run_line_magic(magic_name='picatrixmagics', line='') 59 | self.assertions.assertFalse(magics[magics.name == 'my_silly_magic'].empty) 60 | 61 | line = ip.run_line_magic( 62 | magic_name='my_silly_magic', line='--magnitude 23 this is my string') 63 | 64 | expected_return = ( 65 | 'This magical magic produced 23 magics of this is my string') 66 | self.assertions.assertEqual(line, expected_return) 67 | 68 | 69 | manager.EndToEndTestManager.register_test(BasicTest) 70 | -------------------------------------------------------------------------------- /end_to_end_tests/interface.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google Inc. All rights reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Interface for end-to-end tests.""" 15 | 16 | import collections 17 | import inspect 18 | import logging 19 | import unittest 20 | from typing import Callable, Generator, Text, Tuple 21 | 22 | from IPython.terminal.interactiveshell import TerminalInteractiveShell 23 | 24 | logger = logging.getLogger('picatrix.e2e_test') 25 | 26 | # Default values based on Docker config. 27 | TEST_DATA_DIR = '/usr/local/src/picatrix/end_to_end_tests/test_data' 28 | 29 | 30 | class BaseEndToEndTest: 31 | """Base class for end to end tests. 32 | 33 | Attributes: 34 | assertions: Instance of unittest.TestCase 35 | """ 36 | assertions: unittest.TestCase 37 | 38 | NAME = 'name' 39 | 40 | def __init__(self): 41 | """Initialize the end-to-end test object.""" 42 | self.assertions = unittest.TestCase() 43 | self._counter = collections.Counter() 44 | 45 | def _get_test_methods( 46 | self 47 | ) -> Generator[Tuple[Text, Callable[[TerminalInteractiveShell], None]], None, 48 | None]: 49 | """Inspect class and list all methods that matches the criteria. 50 | 51 | Yields: 52 | Function name and bound method. 53 | """ 54 | for name, func in inspect.getmembers(self, predicate=inspect.ismethod): 55 | if name.startswith('test_'): 56 | yield name, func 57 | 58 | def setup(self): 59 | """Setup function that is run before any tests. 60 | 61 | This is a good place to import any data that is needed. 62 | """ 63 | 64 | def run_tests(self, ip: TerminalInteractiveShell) -> collections.Counter: 65 | """Run all test functions from the class. 66 | 67 | Returns: 68 | Counter of number of tests and errors. 69 | """ 70 | logger.info('*** %s ***', self.NAME) 71 | for test_name, test_func in self._get_test_methods(): 72 | self._counter['tests'] += 1 73 | logger.info('Running test: %s ...', test_name) 74 | try: 75 | test_func(ip) 76 | except Exception: # pylint: disable=broad-except 77 | logger.error('Error while running test %s', self.NAME, exc_info=True) 78 | self._counter['errors'] += 1 79 | continue 80 | logger.info('%s [OK]', test_name) 81 | return self._counter 82 | -------------------------------------------------------------------------------- /end_to_end_tests/manager.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google Inc. All rights reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """This file contains a class for managing end to end tests.""" 15 | from typing import Generator, Text, Tuple 16 | 17 | from . import interface 18 | 19 | 20 | class EndToEndTestManager: 21 | """The test manager.""" 22 | 23 | _class_registry: dict = {} 24 | _exclude_registry: set = set() 25 | 26 | @classmethod 27 | def get_tests( 28 | cls) -> Generator[Tuple[Text, interface.BaseEndToEndTest], None, None]: 29 | """Retrieves the registered tests. 30 | 31 | Yields: 32 | tuple: containing: 33 | str: the uniquely identifying name of the test 34 | type: the test class. 35 | """ 36 | for test_name, test_class in iter(cls._class_registry.items()): 37 | if test_name in cls._exclude_registry: 38 | continue 39 | yield test_name, test_class 40 | 41 | @classmethod 42 | def get_test(cls, test_name: Text) -> interface.BaseEndToEndTest: 43 | """Retrieves a class object of a specific test. 44 | 45 | Args: 46 | test_name (str): name of the test to retrieve. 47 | 48 | Returns: 49 | Instance of Test class object. 50 | 51 | Raises: 52 | KeyError: if the test is not registered. 53 | """ 54 | try: 55 | test_class = cls._class_registry[test_name.lower()] 56 | except KeyError as exc: 57 | raise KeyError( 58 | 'No such test type: {0:s}'.format(test_name.lower())) from exc 59 | return test_class 60 | 61 | @classmethod 62 | def register_test( 63 | cls, 64 | test_class: interface.BaseEndToEndTest, 65 | exclude_from_list: bool = False): 66 | """Registers an test class. 67 | 68 | The test classes are identified by their lower case name. 69 | 70 | Args: 71 | test_class (type): the test class to register. 72 | exclude_from_list (boolean): if set to True then the test 73 | gets registered but will not be included in the 74 | get_tests function. Defaults to False. 75 | 76 | Raises: 77 | KeyError: if class is already set for the corresponding name. 78 | """ 79 | test_name = test_class.NAME.lower() 80 | if test_name in cls._class_registry: 81 | raise KeyError( 82 | 'Class already set for name: {0:s}.'.format(test_class.NAME)) 83 | cls._class_registry[test_name] = test_class 84 | if exclude_from_list: 85 | cls._exclude_registry.add(test_name) 86 | 87 | @classmethod 88 | def clear_registration(cls): 89 | """Clears all test registrations.""" 90 | cls._class_registry = {} 91 | cls._exclude_registry = set() 92 | -------------------------------------------------------------------------------- /end_to_end_tests/timesketch.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google Inc. All rights reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """End to end tests of common picatrix magics.""" 15 | 16 | from IPython.terminal.interactiveshell import TerminalInteractiveShell 17 | from timesketch_api_client import sketch 18 | 19 | from . import interface, manager 20 | 21 | 22 | class TimesketchTest(interface.BaseEndToEndTest): 23 | """End to end tests for Timesketch magics functionality.""" 24 | 25 | NAME = 'timesketch_test' 26 | 27 | def _setup_client(self, ip: TerminalInteractiveShell): 28 | """Setup the TimesketchAPI object into the IPython session.""" 29 | ip.run_cell(raw_cell='from timesketch_api_client import client') 30 | res = ip.run_cell( 31 | raw_cell=( 32 | '_client = client.TimesketchApi(\n' 33 | ' host_uri="https://demo.timesketch.org",\n' 34 | ' username="demo",\n' 35 | ' password="demo",\n' 36 | ' verify=True,\n' 37 | ' auth_mode="userpass")')) 38 | self.assertions.assertTrue(res.success) 39 | res = ip.run_cell( 40 | raw_cell=( 41 | 'from picatrix.lib import state\n' 42 | 'state_obj = state.state()\n' 43 | 'state_obj.add_to_cache(\'timesketch_client\', _client)\n')) 44 | self.assertions.assertTrue(res.success) 45 | 46 | def _get_sketch(self, ip: TerminalInteractiveShell) -> sketch.Sketch: 47 | """Return a sketch object.""" 48 | self._setup_client(ip) 49 | ip.run_line_magic(magic_name='timesketch_set_active_sketch', line='6') 50 | return ip.run_line_magic(magic_name='timesketch_get_sketch', line='') 51 | 52 | def test_get_sketch(self, ip: TerminalInteractiveShell): 53 | """Test fetching a sketch.""" 54 | sketch_obj = self._get_sketch(ip) 55 | self.assertions.assertEqual(sketch_obj.id, 6) 56 | self.assertions.assertEqual(sketch_obj.name, 'Szechuan Sauce - Challenge') 57 | 58 | def test_list_saved_searches(self, ip: TerminalInteractiveShell): 59 | """Test listing up the available saved searches for a sketch.""" 60 | _ = self._get_sketch(ip) 61 | views = ip.run_line_magic( 62 | magic_name='timesketch_list_saved_searches', line='') 63 | expected_views = set( 64 | [ 65 | '18:Szechuan Hits', 66 | '19:Szechuan All Hits', 67 | '16:email_addresses', 68 | '128:Wifitask', 69 | '140:Windows Crash activity', 70 | '139:SSH session view', 71 | '138:Sigma Rule matches', 72 | ]) 73 | self.assertions.assertEqual(set(views.keys()), expected_views) 74 | 75 | def test_query_data(self, ip: TerminalInteractiveShell): 76 | """Test querying for data in a sketch.""" 77 | _ = self._get_sketch(ip) 78 | search_obj = ip.run_line_magic( 79 | magic_name='timesketch_query', 80 | line=( 81 | '--fields datetime,origin,message,hostname,name secret AND ' 82 | 'data_type:"windows:shell_item:file_entry"')) 83 | df = search_obj.table 84 | df_slice = df[df.origin == 'Beth_Secret.lnk'] 85 | self.assertions.assertTrue(df_slice.shape[0] > 0) 86 | origin_set = set(df.origin.unique()) 87 | expected_set = set( 88 | [ 89 | '9b9cdc69c1c24e2b.automaticDestinations-ms', 'Beth_Secret.lnk', 90 | 'HKEY_CURRENT_USER\\Software\\Classes\\Local Settings\\Software' 91 | '\\Microsoft\\Windows\\Shell\\BagMRU\\0\\0\\0', 'NoJerry.lnk', 92 | 'PortalGunPlans.lnk', 'SECRET_beth.lnk', 'Secret.lnk', 93 | 'Szechuan Sauce.lnk', 'f01b4d95cf55d32a.automaticDestinations-ms' 94 | ]) 95 | 96 | self.assertions.assertSetEqual(origin_set, expected_set) 97 | 98 | def test_context_date(self, ip: TerminalInteractiveShell): 99 | """Test querying for contex surrounding a date.""" 100 | _ = self._get_sketch(ip) 101 | search_obj = ip.run_line_magic( 102 | magic_name='timesketch_context_date', 103 | line=( 104 | '--minutes 10 --fields datetime,message,data_type,event_identifier' 105 | ',username,workstation 2020-09-18T22:24:36')) 106 | df = search_obj.table 107 | 108 | self.assertions.assertTrue(df.shape[0] > 5000) 109 | logon_df = df[df.event_identifier == 4624] 110 | logged_in_users = list(logon_df.username.unique()) 111 | self.assertions.assertTrue('Administrator' in logged_in_users) 112 | self.assertions.assertTrue('DWM-1' in logged_in_users) 113 | 114 | df.sort_values('datetime', inplace=True) 115 | first_series = df.iloc[0] 116 | last_series = df.iloc[-1] 117 | 118 | first_time = first_series.datetime 119 | last_time = last_series.datetime 120 | delta = last_time - first_time 121 | delta_rounded = delta.round('min') 122 | 123 | # This should be 10 minutes or 600 seconds. 124 | self.assertions.assertTrue(delta_rounded.total_seconds() == 600.0) 125 | 126 | 127 | manager.EndToEndTestManager.register_test(TimesketchTest) 128 | -------------------------------------------------------------------------------- /end_to_end_tests/tools/run_in_container.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google Inc. All rights reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Script to run all end to end tests.""" 15 | 16 | import time 17 | from collections import Counter 18 | 19 | from IPython.testing.globalipapp import get_ipython 20 | 21 | from end_to_end_tests import manager as test_manager 22 | 23 | manager = test_manager.EndToEndTestManager() 24 | counter = Counter() 25 | 26 | if __name__ == '__main__': 27 | # Sleep to make sure all containers are operational 28 | time.sleep(30) # seconds 29 | ip = get_ipython() 30 | ip.run_cell(raw_cell='from picatrix import notebook_init') 31 | ip.run_cell(raw_cell='notebook_init.init()') 32 | 33 | for name, cls in manager.get_tests(): 34 | test_class = cls() 35 | # Prepare the test environment. 36 | test_class.setup() 37 | # Run all tests. 38 | run_counter = test_class.run_tests(ip) 39 | counter['tests'] += run_counter['tests'] 40 | counter['errors'] += run_counter['errors'] 41 | 42 | successful_tests = counter['tests'] - counter['errors'] 43 | print( 44 | '{0:d} total tests: {1:d} successful and {2:d} failed'.format( 45 | counter['tests'], successful_tests, counter['errors'])) 46 | -------------------------------------------------------------------------------- /install.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # This is a very simple installation script that will only 3 | # work on Linux based systems. 4 | # This will create a data folder and make sure the permissions 5 | # are correct. 6 | 7 | mkdir -p ${HOME}/picadata 8 | cp -r notebooks ${HOME}/picadata/example_notebooks 9 | sudo chgrp -R 1000 ${HOME}/picadata 10 | find ${HOME}/picadata -type d -exec chmod 770 {} \; 11 | find ${HOME}/picadata -type f -exec chmod 660 {} \; 12 | 13 | cd docker 14 | cat docker-compose.yml| sed -e 's/\/tmp\//~\/picadata/g' > docker-tmp.yml 15 | sudo docker-compose -f docker-tmp.yml up -d 16 | cd .. 17 | 18 | echo "Open http://localhost:8899/?token=picatrix in a browser window." 19 | -------------------------------------------------------------------------------- /notebooks/adding_magic.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "nbformat": 4, 3 | "nbformat_minor": 0, 4 | "metadata": { 5 | "colab": { 6 | "name": "adding_magic.ipynb", 7 | "provenance": [], 8 | "private_outputs": true, 9 | "authorship_tag": "ABX9TyMNqvpugAPhDO8vhEebjuY4", 10 | "include_colab_link": true 11 | }, 12 | "kernelspec": { 13 | "name": "python3", 14 | "display_name": "Python 3" 15 | } 16 | }, 17 | "cells": [ 18 | { 19 | "cell_type": "markdown", 20 | "metadata": { 21 | "id": "view-in-github", 22 | "colab_type": "text" 23 | }, 24 | "source": [ 25 | "\"Open" 26 | ] 27 | }, 28 | { 29 | "cell_type": "markdown", 30 | "metadata": { 31 | "id": "gggtpdmgZzi7" 32 | }, 33 | "source": [ 34 | "# Adding A Magic\n", 35 | "\n", 36 | "This notebook describes how to add a magic or register a function into the picatrix set of magics.\n", 37 | "\n", 38 | "## Import\n", 39 | "\n", 40 | "The first thing to do is install the picatrix framework and then import the libraries\n", 41 | "\n", 42 | "(only need to install if you are running a colab hosted kernel)" 43 | ] 44 | }, 45 | { 46 | "cell_type": "code", 47 | "metadata": { 48 | "id": "1ZmkwahOaVa6" 49 | }, 50 | "source": [ 51 | "#@title Only execute if you are connecting to a hosted kernel\n", 52 | "!pip install picatrix" 53 | ], 54 | "execution_count": null, 55 | "outputs": [] 56 | }, 57 | { 58 | "cell_type": "code", 59 | "metadata": { 60 | "id": "QZdMGQAuZwxp" 61 | }, 62 | "source": [ 63 | "from picatrix.lib import framework\n", 64 | "from picatrix.lib import utils\n", 65 | "\n", 66 | "# This should not be included in the magic definition file, only used\n", 67 | "# in this notebook since we are comparing all magic registration.\n", 68 | "from picatrix import notebook_init\n", 69 | "\n", 70 | "notebook_init.init()" 71 | ], 72 | "execution_count": null, 73 | "outputs": [] 74 | }, 75 | { 76 | "cell_type": "markdown", 77 | "metadata": { 78 | "id": "1kh5DfHzaDxt" 79 | }, 80 | "source": [ 81 | "Then we need to create a function:" 82 | ] 83 | }, 84 | { 85 | "cell_type": "code", 86 | "metadata": { 87 | "id": "OeOmXWPhaGKW" 88 | }, 89 | "source": [ 90 | "from typing import Optional\n", 91 | "from typing import Text\n", 92 | "\n", 93 | "@framework.picatrix_magic\n", 94 | "def my_silly_magic(data: Text, magnitude: Optional[int] = 100) -> Text:\n", 95 | " \"\"\"Return a silly string with no meaningful value.\n", 96 | "\n", 97 | " Args:\n", 98 | " data (str): This is a string that will be printed back.\n", 99 | " magnitude (int): A number that will be displayed in the string.\n", 100 | "\n", 101 | " Returns:\n", 102 | " A string that basically combines the two options.\n", 103 | " \"\"\" \n", 104 | " return f'This magical magic produced {magnitude} magics of {data.strip()}'" 105 | ], 106 | "execution_count": null, 107 | "outputs": [] 108 | }, 109 | { 110 | "cell_type": "markdown", 111 | "metadata": { 112 | "id": "AXdg1ah-ayBH" 113 | }, 114 | "source": [ 115 | "In order to register a magic it has to have few properties:\n", 116 | "\n", 117 | "1. Be a regular Python function that accepts parameters (optional if it returns a value)\n", 118 | "2. The first argument it must accept is `data` (this is due to how magics work). If you don't need an argument, set the default value of `data` to an empty string.\n", 119 | "3. Use typing to denote the type of the argument values.\n", 120 | "4. The function must include a docstring, where the first line describes the function.\n", 121 | "5. The docstring also must have an argument section, where each argument is further described (this is used to generate the helpstring for the magic/function).\n", 122 | "6. If the function returns a value it must define a Returns section.\n", 123 | "\n", 124 | "Once these requirements are fulfilled, a simple decorator is all that is required to register the magic and make sure it is available.\n", 125 | "\n", 126 | "## Test the Magic\n", 127 | "\n", 128 | "Now once the magic has been registered we can first test to see if it is registered:" 129 | ] 130 | }, 131 | { 132 | "cell_type": "code", 133 | "metadata": { 134 | "id": "w-CqWQyCTj3t" 135 | }, 136 | "source": [ 137 | "%picatrixmagics" 138 | ], 139 | "execution_count": null, 140 | "outputs": [] 141 | }, 142 | { 143 | "cell_type": "markdown", 144 | "metadata": { 145 | "id": "0eEIt9KjTk_V" 146 | }, 147 | "source": [ 148 | "This does produce quite a lot of values, let's filter it out:" 149 | ] 150 | }, 151 | { 152 | "cell_type": "code", 153 | "metadata": { 154 | "id": "2QAcImyzTm_C" 155 | }, 156 | "source": [ 157 | "magics = %picatrixmagics\n", 158 | "\n", 159 | "magics[magics.name.str.contains('silly_magic')]" 160 | ], 161 | "execution_count": null, 162 | "outputs": [] 163 | }, 164 | { 165 | "cell_type": "markdown", 166 | "metadata": { 167 | "id": "IQ2kVmB4UFPE" 168 | }, 169 | "source": [ 170 | "OK, we can see that it is registered. Now let's try to call it:" 171 | ] 172 | }, 173 | { 174 | "cell_type": "code", 175 | "metadata": { 176 | "id": "b87RuDZAah8G" 177 | }, 178 | "source": [ 179 | "%my_silly_magic foobar" 180 | ], 181 | "execution_count": null, 182 | "outputs": [] 183 | }, 184 | { 185 | "cell_type": "markdown", 186 | "metadata": { 187 | "id": "UH9WIGKSUIye" 188 | }, 189 | "source": [ 190 | "And check out it's help message:" 191 | ] 192 | }, 193 | { 194 | "cell_type": "code", 195 | "metadata": { 196 | "id": "VVlg7pzfUKIo" 197 | }, 198 | "source": [ 199 | "%my_silly_magic --help" 200 | ], 201 | "execution_count": null, 202 | "outputs": [] 203 | }, 204 | { 205 | "cell_type": "markdown", 206 | "metadata": { 207 | "id": "GmSQRYmhUMXi" 208 | }, 209 | "source": [ 210 | "Here you can see the results from the docstring being used to generate the help for the magic.\n", 211 | "\n", 212 | "Now use the call magic:" 213 | ] 214 | }, 215 | { 216 | "cell_type": "code", 217 | "metadata": { 218 | "id": "SRvCRD6Sa3gN" 219 | }, 220 | "source": [ 221 | "%%my_silly_magic \n", 222 | "this is some text\n", 223 | "and some more text\n", 224 | "and yet even more" 225 | ], 226 | "execution_count": null, 227 | "outputs": [] 228 | }, 229 | { 230 | "cell_type": "markdown", 231 | "metadata": { 232 | "id": "qx0GemDwUT9l" 233 | }, 234 | "source": [ 235 | "And set the arguments:" 236 | ] 237 | }, 238 | { 239 | "cell_type": "code", 240 | "metadata": { 241 | "id": "dqMfcC90a6zC" 242 | }, 243 | "source": [ 244 | "%%my_silly_magic --magnitude 234 store_here\n", 245 | "and here is the text" 246 | ], 247 | "execution_count": null, 248 | "outputs": [] 249 | }, 250 | { 251 | "cell_type": "code", 252 | "metadata": { 253 | "id": "NhvbYDmAa-46" 254 | }, 255 | "source": [ 256 | "store_here" 257 | ], 258 | "execution_count": null, 259 | "outputs": [] 260 | }, 261 | { 262 | "cell_type": "markdown", 263 | "metadata": { 264 | "id": "yGQnPiNzUXZo" 265 | }, 266 | "source": [ 267 | "And finally we can use the exposed function:" 268 | ] 269 | }, 270 | { 271 | "cell_type": "code", 272 | "metadata": { 273 | "id": "R1GQp1TXa_1y" 274 | }, 275 | "source": [ 276 | "my_silly_magic_func?" 277 | ], 278 | "execution_count": null, 279 | "outputs": [] 280 | }, 281 | { 282 | "cell_type": "code", 283 | "metadata": { 284 | "id": "a5HobRkoUak7" 285 | }, 286 | "source": [ 287 | "my_silly_magic_func('some random string', magnitude=234)" 288 | ], 289 | "execution_count": null, 290 | "outputs": [] 291 | }, 292 | { 293 | "cell_type": "code", 294 | "metadata": { 295 | "id": "rIgdH_g5UeMM" 296 | }, 297 | "source": [ 298 | "" 299 | ], 300 | "execution_count": null, 301 | "outputs": [] 302 | } 303 | ] 304 | } 305 | -------------------------------------------------------------------------------- /picatrix/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Sets up Picatrix environment.""" 15 | 16 | from typing import Optional, Text, Tuple 17 | 18 | from .lib.namespace import ( 19 | FeatureContext, 20 | FeatureNamespace, 21 | Function, 22 | RootContext, 23 | RootNamespace, 24 | ) 25 | 26 | px = RootNamespace( 27 | "px", """Picatrix root namespace. 28 | 29 | Features can be accessed as `px.{feature_name}`, e.g. `px.timesketch`. 30 | 31 | Try: 32 | * `px.search(\"keyword\")` to search features by keyword. 33 | * `px?` or `px.{feature_name}?` for help. 34 | """) 35 | ctx = RootContext( 36 | "ctx", """Picatrix root context, holds all runtime parameters. 37 | 38 | To check runtime parameters of a certain feature, try `ctx.{feature_name}`, 39 | e.g. `ctx.timesketch.to_frame(with_values=True). 40 | 41 | Try: 42 | * `ctx.search(\"keyword\")` to search parameters by keyword 43 | * `ctx?` or `ctx.{feature_name}?` for help. 44 | """) 45 | 46 | 47 | def new_namespace( 48 | name: Text, 49 | namespace_docstring: Optional[Text] = None, 50 | context_docstring: Optional[Text] = None, 51 | ) -> Tuple[FeatureNamespace, FeatureContext]: 52 | """Creates new Picatrix namespace and context for the feature. 53 | 54 | Newly registered feature will be available as: 55 | * `px.{name}` for functionalities, 56 | * `ctx.{name}` for runtime parameters, i.e. context. 57 | 58 | Args: 59 | name: new namespaces name 60 | namespace_docstring: custom docstring for the namespace 61 | if None, "`px.{name}` contains all Picatrix features of {name}" 62 | context_docstring: custom docstring for the context, 63 | if None, "`ctx.{name}` contains all Picatrix context for {name}" 64 | 65 | Returns: 66 | Tuple[FeatureNamespace, FeatureContext]: new namespace and context. 67 | 68 | Raises: 69 | NamespaceKeyExistsError: when namespace already exists 70 | NamespaceKeyError: when name is invalid, e.g. isn't Python identifier 71 | """ 72 | return ( 73 | px.add_namespace(name, docstring=namespace_docstring), 74 | ctx.add_namespace(name, docstring=context_docstring)) 75 | 76 | 77 | def new_line_magic(func: Function, name: Optional[Text] = None): 78 | """Adds new line magic to Picatrix namespace. 79 | 80 | Newly registered magic will be available as: 81 | * `%{name}` IPython line magic 82 | * `px.magic.{name}` function 83 | 84 | Args: 85 | func: function to be made into magic 86 | name: optional custom magic name, if None function name will be used 87 | 88 | Raises: 89 | NamespaceKeyExistsError: when required key already exists 90 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 91 | MagicParsingError: when provided function is an invalid magic 92 | """ 93 | px.add_line_magic(func, name) 94 | 95 | 96 | def new_cell_magic(func: Function, name: Optional[Text] = None): 97 | """Adds new cell magic to Picatrix namespace. 98 | 99 | Newly registered magic will be available as: 100 | * `%%{name}` IPython cell magic 101 | * `px.magic.{name}` function 102 | 103 | Args: 104 | func: function to be made into magic 105 | name: optional custom magic name, if None function name will be used 106 | 107 | Raises: 108 | NamespaceKeyExistsError: when required key already exists 109 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 110 | MagicParsingError: when provided function is an invalid magic 111 | """ 112 | px.add_cell_magic(func, name if name else func.__name__) 113 | 114 | 115 | # shouldn't be exported 116 | del Optional, Text, Tuple # type: ignore 117 | del FeatureContext, FeatureNamespace, Function, RootContext, RootNamespace, 118 | -------------------------------------------------------------------------------- /picatrix/helpers/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Class that loads up all helper libraries.""" 15 | 16 | from picatrix.helpers import table 17 | -------------------------------------------------------------------------------- /picatrix/helpers/table.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Defines helper functions to display tables or dataframes.""" 15 | from typing import Optional, Tuple, Union 16 | 17 | import ipyaggrid 18 | import pandas 19 | 20 | try: 21 | # 3rd party widgets do not work inside colab, we will use the 22 | # built-in data table instead there. 23 | from google.colab.data_table import DataTable 24 | except ImportError: 25 | DataTable = None 26 | 27 | from picatrix.lib import framework 28 | 29 | DEFAULT_COLUMNS_HIDE = ('_type', '_id', '__ts_emojis', '_index') 30 | 31 | 32 | @framework.picatrix_helper 33 | def display_table( 34 | data_frame: pandas.DataFrame, 35 | hide_columns: Optional[Tuple[str]] = None 36 | ) -> Union[ipyaggrid.grid.Grid, DataTable]: 37 | """Display a dataframe interactively with a toolbar.""" 38 | if DataTable: 39 | return DataTable(data_frame, include_index=False) 40 | 41 | column_defs = [] 42 | 43 | if hide_columns is None: 44 | hide_columns = DEFAULT_COLUMNS_HIDE 45 | 46 | for column in data_frame.columns: 47 | pivot_group = column != 'message' 48 | hide = False 49 | if column in hide_columns: 50 | hide = True 51 | pivot_group = False 52 | 53 | column_dict = { 54 | 'headerName': column.title(), 55 | 'field': column, 56 | 'rowGroup': False, 57 | 'enableRowGroup': True, 58 | 'hide': hide, 59 | 'pivot': pivot_group, 60 | 'sortable': True, 61 | 'resizable': True, 62 | } 63 | 64 | column_defs.append(column_dict) 65 | 66 | grid_options = { 67 | 'columnDefs': column_defs, 68 | 'enableSorting': True, 69 | 'enableFilter': True, 70 | 'enableColResize': True, 71 | 'enableRangeSelection': True, 72 | 'editable': False, 73 | 'rowGroupPanelShow': 'always', 74 | 'rowSelection': 'multiple', 75 | } 76 | 77 | return ipyaggrid.Grid( 78 | grid_data=data_frame, 79 | quick_filter=True, 80 | show_toggle_edit=False, 81 | export_csv=True, 82 | export_excel=False, 83 | export_to_df=False, 84 | theme='ag-theme-balham', 85 | show_toggle_delete=False, 86 | columns_fit='auto', 87 | index=False, 88 | grid_options=grid_options, 89 | keep_multiindex=True, 90 | ) 91 | -------------------------------------------------------------------------------- /picatrix/lib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/picatrix/b2aa91737273e36526129b69f60e5fc3a8c6fa4e/picatrix/lib/__init__.py -------------------------------------------------------------------------------- /picatrix/lib/error.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Class that defines custom errors for picatrix.""" 15 | 16 | 17 | class Error(Exception): 18 | """Base error class.""" 19 | 20 | 21 | class ArgParserNonZeroStatus(Error): 22 | """Raised when the argument parser has exited with non-zero status.""" 23 | -------------------------------------------------------------------------------- /picatrix/lib/framework.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Class that defines the framework for exported magics and functions.""" 15 | 16 | import argparse 17 | import functools 18 | import inspect 19 | import types 20 | import typing 21 | from typing import Any, Callable, Dict, List, Optional, Text 22 | 23 | try: 24 | # Got introduced in python 3.8. 25 | from typing import Protocol 26 | except ImportError: 27 | from typing_extensions import Protocol 28 | 29 | from picatrix.lib import error, manager, state, utils 30 | 31 | 32 | class MagicProtocol(Protocol): 33 | """Simple Typing protocol class for magic functions.""" 34 | 35 | def __call__(self, *args: Any) -> Any: 36 | pass 37 | 38 | 39 | class MagicArgument: 40 | """Simple Argument holder for magic arguments.""" 41 | 42 | def __init__(self, *argc, **kwargs): 43 | self.argc = argc 44 | self.kwargs = kwargs 45 | 46 | 47 | class MagicArgumentParser(argparse.ArgumentParser): 48 | """Argument parser for picatrix magics.""" 49 | 50 | def __init__(self, *args, **kwargs): 51 | super(MagicArgumentParser, self).__init__(*args, **kwargs) 52 | self.storage = {} 53 | 54 | def exit(self, status: Optional[int] = 0, message: Optional[Text] = ''): 55 | """Exiting method for argument parser. 56 | 57 | Args: 58 | status (int): exit status of the parser. 59 | message (str): the error message. 60 | 61 | Raises: 62 | KeyError: when the parser is unable to parse the arguments. 63 | error.ArgParserNonZeroStatus: when the parser has successfully completed. 64 | """ 65 | if not status: 66 | raise error.ArgParserNonZeroStatus('Exiting.') 67 | 68 | if message: 69 | raise KeyError('Wrong usage: {0:s}'.format(message.strip())) 70 | 71 | raise KeyError('Wrong usage, no further error message supplied.') 72 | 73 | 74 | class _Magic: 75 | """The Picatrix Magic decorator.""" 76 | 77 | argument_parser: MagicArgumentParser 78 | magic_name: Text 79 | 80 | def __init__( 81 | self, 82 | fn: MagicProtocol, 83 | arguments: Optional[List[MagicArgument]] = None, 84 | name_func: Optional[Callable[[Text], Text]] = None): 85 | """Initialize the Picatrix Magic.""" 86 | self.fn = fn 87 | if name_func: 88 | self.magic_name = name_func(fn.__name__) 89 | else: 90 | self.magic_name = fn.__name__ 91 | 92 | functools.update_wrapper(self, fn) 93 | self.__name__ = self.magic_name 94 | 95 | # pylint: disable=access-member-before-definition 96 | first_line = self.__doc__.split('\n')[0] 97 | self.argument_parser = MagicArgumentParser( 98 | prog=self.magic_name, 99 | description=first_line, 100 | usage=( 101 | '%%%(prog)s [arguments] data\nor\n%%%%%(prog)s [arguments]\ndata')) 102 | 103 | if arguments: 104 | for argument in arguments: 105 | dest_name = argument.kwargs.get('dest') 106 | type_string = argument.kwargs.get('type') 107 | # We need to check whether the custom argument contains an object 108 | # type, if so we will need to change it into a str and store the 109 | # type in the argument parser so it can be confirmed after variable 110 | # expansion. 111 | if dest_name and type_string: 112 | if isinstance(type_string, type): 113 | type_string = str(type_string).split('\'')[1] 114 | if type_string not in ('str', 'int', 'bool', 'float', 'unicode'): 115 | self.argument_parser.storage[dest_name] = type_string 116 | argument.kwargs['type'] = str 117 | 118 | self.argument_parser.add_argument(*argument.argc, **argument.kwargs) 119 | else: 120 | _add_function_arguments_to_parser(fn, self.argument_parser) 121 | 122 | try: 123 | self.argument_parser.add_argument( 124 | '--bindto', 125 | '-bindto', 126 | type=str, 127 | default='', 128 | action='store', 129 | dest='bindto', 130 | help='Bind the results to a variable instead of being returned.') 131 | except argparse.ArgumentError: 132 | pass 133 | 134 | try: 135 | self.argument_parser.add_argument( 136 | 'data', 137 | nargs='?', 138 | help=self.argument_parser.storage.get('_data_help', '')) 139 | except argparse.ArgumentError: 140 | pass 141 | 142 | self.__doc__ = self.argument_parser.format_help() 143 | 144 | # pylint: disable=inconsistent-return-statements 145 | def __call__(self, line: Text, cell: Optional[Text] = None) -> Optional[Any]: 146 | line_magic = cell is None 147 | 148 | arguments = _parse_line_string(line) 149 | 150 | try: 151 | options = self.argument_parser.parse_args(arguments) 152 | except KeyError as e: 153 | print( 154 | ( 155 | 'Unable to parse arguments, with error: {0!s}.\n' 156 | 'Correct usage is: {1:s}').format( 157 | e, self.argument_parser.format_help())) 158 | return 159 | except error.ArgParserNonZeroStatus: 160 | # When argparser ends execution but without an error this exception 161 | # is raised, eg: when "-h" or help is used. In those cases we need 162 | # to return without running the magic function. 163 | return 164 | 165 | if not line_magic: 166 | variable = options.data 167 | options.data = cell 168 | 169 | bind_to = options.bindto 170 | option_dict = options.__dict__ 171 | del option_dict['bindto'] 172 | 173 | # TODO: Change this so that there is no variable expansion 174 | # done by the core and all variable expansion is only done here. 175 | for key, value in iter(option_dict.items()): 176 | if not value: 177 | continue 178 | 179 | if not isinstance(value, str): 180 | continue 181 | 182 | if value[0] == '{' and value[-1] == '}': 183 | var_name = value[1:-1] 184 | var_type = self.argument_parser.storage.get(key) 185 | 186 | if '{' not in var_name and '}' not in var_name: 187 | var_obj = utils.ipython_get_global(var_name) 188 | if var_type and var_type != 'object': 189 | type_string = type(var_obj).__name__ 190 | if (type_string != var_type) and not type_string.endswith( 191 | var_type) and not var_type.endswith(type_string): 192 | raise KeyError( 193 | ( 194 | 'Variable [{0:s}] is not of the correct type [{1:s}] for ' 195 | 'this magic. Type is: {2!s}').format( 196 | var_name, var_type, type(var_obj))) 197 | option_dict[key] = var_obj 198 | 199 | return_value = self.fn(**option_dict) 200 | state_obj = state.state() 201 | 202 | if not line_magic and variable: 203 | bind_to = variable 204 | 205 | return state_obj.set_output( 206 | output=return_value, magic_name=self.magic_name, bind_to=bind_to) 207 | 208 | def __dir__(self): 209 | options = getattr(self.argument_parser, '_option_string_actions', {}) 210 | return list(options.keys()) 211 | 212 | def __getattribute__(self, name: Text): 213 | # Overwriting function to behave like a function when called 214 | # by isinstance, in order to use the inspect module as well as to 215 | # produce a better help message. 216 | if name == '__class__': 217 | return types.FunctionType 218 | 219 | if name.startswith('func'): 220 | return self.fn.__getattribute__(name) # pytype: disable=attribute-error 221 | 222 | return super(_Magic, self).__getattribute__(name) 223 | 224 | 225 | def _get_arguments_from_arg_lines( 226 | arg_lines: List[Text]) -> List[Dict[Text, Text]]: 227 | """Return a list of parsed arguments from argument docstring. 228 | 229 | Args: 230 | arg_lines (list): a list of lines, one per argument as defined in 231 | the docstring. 232 | 233 | Returns: 234 | A list of dict, one per argument. Each dict will have the following 235 | attributes; variable, type and description. 236 | """ 237 | args = [] 238 | for arg_line in arg_lines: 239 | var_string, _, description = arg_line.partition(':') 240 | var, var_type = var_string.split() 241 | argument = { 242 | 'variable': var, 243 | 'type': var_type[1:-1], 244 | 'description': description.strip() 245 | } 246 | args.append(argument) 247 | return args 248 | 249 | 250 | def _get_argument_lines_from_docstring(lines: List[Text]) -> List[Text]: 251 | """Return names and types of arguments read from a doc string. 252 | 253 | This function takes in a function's docstring and returns back 254 | the argument section (Args:), with one line per argument defined. 255 | Arguments can be multi-lines. 256 | 257 | Args: 258 | lines (list): a list of lines read from a function's docstring. 259 | 260 | Returns: 261 | A list of lines that contain argument definitions. 262 | """ 263 | arg_lines = [] 264 | 265 | try: 266 | arg_index = [x.strip() for x in lines].index('Args:') 267 | except ValueError: 268 | return arg_lines 269 | 270 | new_index = arg_index + 1 271 | spaces = lines[arg_index].index('A') 272 | while True: 273 | line = lines[new_index] 274 | 275 | if not line.strip(): 276 | break 277 | 278 | if line.strip() in ('Returns:', 'Raises:'): 279 | break 280 | 281 | # Checking indentation, whether the argument line is continuing or 282 | # a new argument is being defined. 283 | space_count = len(line) - len(line.lstrip(' ')) 284 | if space_count == spaces * 2: 285 | arg_lines.append(line.strip()) 286 | else: 287 | arg_lines[-1] = '{0:s} {1:s}'.format(arg_lines[-1], line.strip()) 288 | 289 | new_index += 1 290 | if new_index >= len(lines) - 1: 291 | break 292 | 293 | return arg_lines 294 | 295 | 296 | def _add_function_arguments_to_parser( 297 | fn: MagicProtocol, parser: MagicArgumentParser): 298 | """Adds arguments to a parser from parsing a function. 299 | 300 | Args: 301 | fn (function): a function to extract an argument parser from. 302 | parser (MagicArgumentParser): an argument parser object, that will be used 303 | to add arguments to. 304 | 305 | Raises: 306 | ValueError: when the docstring is not correctly formatted. 307 | """ 308 | try: 309 | doc_string_args = _parse_docstring(fn) 310 | except ValueError as e: 311 | raise ValueError( 312 | 'Unable to register the magic %{0:s} since the docstring ' 313 | 'of the function is not correctly formatted. Error message ' 314 | 'is: {1!s}'.format(fn.__name__, e)) from e 315 | 316 | inspect_fn = inspect.getfullargspec 317 | 318 | try: 319 | spec = inspect_fn(fn) # pytype: disable=wrong-arg-types 320 | except TypeError: 321 | spec = None 322 | 323 | if spec and spec.defaults: 324 | defaults_len = len(spec.defaults) 325 | default_values = dict(zip(spec.args[-defaults_len:], spec.defaults)) 326 | else: 327 | default_values = {} 328 | 329 | for argument in doc_string_args: 330 | arg_type_string = argument.get('type', 'str') 331 | description = argument.get('description', '') 332 | # Percent sign has a special meaning for argument parser, escaping potential 333 | # percent signs in docstrings. 334 | description = description.replace('%', '%%') 335 | 336 | if argument.get('variable') == 'data': 337 | parser.storage['data'] = arg_type_string 338 | parser.storage['_data_help'] = description 339 | continue 340 | 341 | variable = argument.get('variable', '') 342 | arg_type = str 343 | if arg_type_string in ('str', 'unicode'): 344 | arg_type = str 345 | elif arg_type_string == 'float': 346 | arg_type = float 347 | elif arg_type_string == 'int': 348 | arg_type = int 349 | elif arg_type_string == 'bool': 350 | arg_type = bool 351 | arg_default = default_values.get(variable) 352 | if arg_default: 353 | action = 'store_false' 354 | else: 355 | action = 'store_true' 356 | else: 357 | # Type is a string reference to an object. 358 | arg_type = str 359 | parser.storage[variable] = arg_type_string 360 | 361 | if arg_type == bool: 362 | parser.add_argument( 363 | '--{0:s}'.format(variable), 364 | '-{0:s}'.format(variable), 365 | dest=variable, 366 | action=action, 367 | default=default_values.get(variable), 368 | help=description) 369 | else: 370 | parser.add_argument( 371 | '--{0:s}'.format(variable), 372 | '-{0:s}'.format(variable), 373 | dest=variable, 374 | type=arg_type, 375 | action='store', 376 | default=default_values.get(variable), 377 | help=description) 378 | 379 | 380 | def _parse_docstring(function: MagicProtocol) -> List[Dict[Text, Text]]: 381 | """Return a list of arguments extracted from a function's docstring. 382 | 383 | Args: 384 | function (function): a function to extract the docstring from. 385 | 386 | Returns: 387 | A list of dict, one per argument. Each dict will have the following 388 | attributes; variable, type and description. 389 | """ 390 | doc_string = getattr(function, '__doc__') 391 | if not doc_string: 392 | return [] 393 | 394 | doc_lines = doc_string.split('\n') 395 | arg_lines = _get_argument_lines_from_docstring(doc_lines) 396 | return _get_arguments_from_arg_lines(arg_lines) 397 | 398 | 399 | def _parse_line_string(line: Text) -> List[Text]: 400 | """Return a list of arguments from a line magic. 401 | 402 | Args: 403 | line (str): the value passed to the line magic, or line attribute in a 404 | cell magic. 405 | 406 | Returns: 407 | List of arguments that can be parsed by argparse. 408 | """ 409 | arguments = [] 410 | temp_args = [] 411 | in_quotes = False 412 | arg_item = False 413 | quota_char = '' 414 | 415 | for item in line.strip().split(): 416 | # Two cases to watch for, line starts with - (argument) or quotes. 417 | if in_quotes: 418 | if item[-1] == quota_char: 419 | temp_args.append(item[:-1]) 420 | arguments.append(' '.join(temp_args)) 421 | temp_args = [] 422 | in_quotes = False 423 | quota_char = '' 424 | else: 425 | temp_args.append(item) 426 | continue 427 | 428 | if item[0] == '-': 429 | arguments.append(item) 430 | arg_item = True 431 | continue 432 | 433 | if item[0] in ('\'', '"'): 434 | quota_char = item[0] 435 | 436 | # The quoted argument is a single word. 437 | if item[-1] == quota_char: 438 | arguments.append(item[1:-1]) 439 | quota_char = '' 440 | continue 441 | 442 | temp_args.append(item[1:]) 443 | in_quotes = True 444 | quota_char = item[0] 445 | continue 446 | 447 | if arg_item: 448 | arguments.append(item) 449 | arg_item = False 450 | else: 451 | temp_args.append(item) 452 | 453 | if temp_args: 454 | arguments.append(' '.join(temp_args)) 455 | 456 | return arguments 457 | 458 | 459 | def picatrix_helper(function: Callable[..., Any]) -> Callable[..., Any]: 460 | """Decorator to register a picatrix helper. 461 | 462 | Args: 463 | function (function): if the decorator is called without any 464 | arguments the helper function is passed to the decorator. 465 | 466 | Returns: 467 | The function that was passed in. 468 | """ 469 | typing_hints = typing.get_type_hints(function) 470 | manager.MagicManager.register_helper( 471 | name=function.__name__, helper=function, typing_help=typing_hints) 472 | 473 | try: 474 | _ = utils.ipython_get_global(function.__name__) 475 | except KeyError: 476 | utils.ipython_bind_global(function.__name__, function) 477 | 478 | return function 479 | 480 | 481 | def picatrix_magic( 482 | function: Optional[MagicProtocol] = None, 483 | arguments: Optional[List[MagicArgument]] = None, 484 | name_func: Optional[Callable[[Text], Text]] = None, 485 | conditional: Optional[Callable[[None], bool]] = None) -> MagicProtocol: 486 | """Decorator to turn functions into IPYthon magics for picatrix. 487 | 488 | Args: 489 | function (function): if the decorator is called without any arguments 490 | the magic function is passed to the decorator. 491 | arguments (list): list of MagicArgument objects to pass to the magic 492 | argument parser. 493 | name_func (function: a name function that will accept a single argument 494 | and return back a name that will be used to register the magic. 495 | Optional and if not provide the name of tbe function will be used. 496 | conditional (function): a function that should return a bool, used to 497 | determine whether to register magic or not. This can be used by 498 | magics to determine whether a magic should be registered or not, for 499 | instance basing that on whether a certain magic is able to reach the 500 | service it requires. This is optional and if not provided a magic 501 | will be registered. 502 | 503 | Returns: 504 | the decorator function. 505 | """ 506 | if function: 507 | magic_function = _Magic(function, name_func=name_func, arguments=arguments) 508 | manager.MagicManager.register_magic(magic_function, conditional) 509 | return magic_function 510 | 511 | def wrapper(func): 512 | """Wrapper for the magic.""" 513 | magic_function = _Magic(func, name_func=name_func, arguments=arguments) 514 | 515 | manager.MagicManager.register_magic(magic_function, conditional) 516 | return magic_function 517 | 518 | return wrapper 519 | -------------------------------------------------------------------------------- /picatrix/lib/framework_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Copyright 2020 Google LLC 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # https://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """Tests for the picatrix framework.""" 16 | from typing import Optional, Text 17 | 18 | from picatrix.lib import framework, manager 19 | 20 | 21 | def my_very_own_test_magic(data: Text, stuff: Optional[int] = 20) -> Text: 22 | """This is a magic that is used for testing. 23 | 24 | Args: 25 | data (str): This is a string. 26 | stuff (int): And this is a number. 27 | 28 | Returns: 29 | str: A string that combines the two parameters together. 30 | """ 31 | return f'{data.strip()} - {stuff}' 32 | 33 | 34 | def test_registration(): 35 | """Test the magic decorator.""" 36 | magic = framework.picatrix_magic(my_very_own_test_magic) 37 | assert magic.__doc__.startswith('usage: %my_very_own_test_magic') 38 | results = magic(line='--stuff 23 this is a text') 39 | assert results == 'this is a text - 23' 40 | 41 | manager.MagicManager.deregister_magic(magic.magic_name) 42 | -------------------------------------------------------------------------------- /picatrix/lib/ipython.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # type: ignore 16 | """Helper function for Picatrix IPython magic handling. 17 | 18 | This module exists mostly because of IPython non-strict type checking which 19 | annoys pyright. Type checking is disabled for this module so tread lightly. 20 | """ 21 | 22 | from typing import Any, Callable, Text, Union 23 | 24 | from IPython import get_ipython 25 | from IPython.core.interactiveshell import InteractiveShell 26 | 27 | 28 | def add_line_magic(name: Text, magic: Callable[..., Any]): 29 | """Adds new IPython line magic (`%magic`).""" 30 | ip: Union[InteractiveShell, None] = get_ipython() 31 | if not ip: 32 | return 33 | ip.register_magic_function(magic, magic_kind="line", magic_name=name) 34 | 35 | 36 | def add_cell_magic(name: Text, magic: Callable[..., Any]): 37 | """Adds new IPython line magic (`%%magic`).""" 38 | ip: Union[InteractiveShell, None] = get_ipython() 39 | if not ip: 40 | return 41 | ip.register_magic_function(magic, magic_kind="cell", magic_name=name) 42 | 43 | 44 | def delete_magic(name: Text): 45 | """Deletes IPython magic based on name.""" 46 | ip: Union[InteractiveShell, None] = get_ipython() 47 | if not ip: 48 | return 49 | 50 | line_magics = ip.magics_manager.magics.get("line", {}) 51 | if name in line_magics: 52 | _ = ip.magics_manager.magics["line"].pop(name) 53 | cell_magics = ip.magics_manager.magics.get("cell", {}) 54 | if name in cell_magics: 55 | _ = ip.magics_manager.magics["cell"].pop(name) 56 | -------------------------------------------------------------------------------- /picatrix/lib/magic.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Types and functions defining Picatrix integration with IPython magics.""" 15 | 16 | import argparse 17 | import shlex 18 | from dataclasses import dataclass, field 19 | from enum import Enum, auto 20 | from inspect import getfullargspec 21 | from typing import ( 22 | Any, 23 | Callable, 24 | Dict, 25 | Iterable, 26 | List, 27 | NoReturn, 28 | Optional, 29 | Text, 30 | Type, 31 | Union, 32 | ) 33 | 34 | from docstring_parser import parse 35 | 36 | from .error import Error 37 | from .utils import ipython_bind_global 38 | 39 | 40 | class MagicError(Error): 41 | """Generic Picatrix error related to IPython magics.""" 42 | 43 | 44 | class MagicParsingError(MagicError, ValueError): 45 | """Raised when invalid function is provided to magic framework.""" 46 | 47 | 48 | class MagicArgParsingError(MagicError, ValueError): 49 | """Raised when invalid arguments where provided to magic.""" 50 | 51 | 52 | class MagicType(Enum): 53 | """Enum indicating a type of the IPython magic.""" 54 | CELL = auto() 55 | LINE = auto() 56 | 57 | 58 | _MagicArgValues = [bool, int, float, str] 59 | MagicArgValue = Union[bool, int, float, str] 60 | _MagicArgValueType = Union[Type[bool], Type[int], Type[float], Type[str]] 61 | 62 | 63 | def _validate_arg_value(arg: Text, v: Any) -> MagicArgValue: 64 | for typ in _MagicArgValues: 65 | if isinstance(v, typ): 66 | return typ(v) 67 | raise MagicArgParsingError( 68 | f"{arg}=\"{v}\" is of invalid type; has to be {MagicArgValue}") 69 | 70 | 71 | @dataclass(frozen=True) 72 | class MagicArgs: 73 | """Arguments parsed out of IPython magic invocation.""" 74 | bind_variable: Text = field(default="_") 75 | kwargs: Dict[Text, MagicArgValue] = field(default_factory=dict) 76 | 77 | def __post_init__(self): 78 | if not self.bind_variable.isidentifier(): # pylint: disable=no-member 79 | raise MagicArgParsingError( 80 | f"\"{self.bind_variable}\" isn't valid Python " 81 | "identifier, see: " 82 | "https://docs.python.org/3/reference/" 83 | "lexical_analysis.html#identifiers") 84 | 85 | 86 | class ArgParserNonZeroStatus(Exception): 87 | """Raised when the argument parser has exited with non-zero status.""" 88 | 89 | 90 | class MagicArgumentParser(argparse.ArgumentParser): 91 | """Argument parser for Picatrix magics.""" 92 | 93 | def exit(self, status: int = 0, message: Optional[Text] = None) -> NoReturn: 94 | """Exiting method for argument parser. 95 | 96 | Args: 97 | status (int): exit status of the parser. 98 | message (str): the error message. 99 | 100 | Raises: 101 | MagicArgParsingError: when the parser is unable to parse the arguments. 102 | ArgParserNonZeroStatus: when the parser has successfully completed. 103 | """ 104 | if not status: 105 | raise ArgParserNonZeroStatus() 106 | 107 | if message: 108 | raise MagicArgParsingError(message.strip()) 109 | 110 | raise MagicArgParsingError("Wrong usage.") 111 | 112 | def parse_magic_args( 113 | self, line: Text, cell: Optional[Text] = None) -> MagicArgs: 114 | """Parse arguments out of line and cell content as provided by IPython.""" 115 | line = line.strip() 116 | 117 | kwargs: Dict[Text, MagicArgValue] = {} 118 | if cell: 119 | kwargs["cell"] = cell 120 | 121 | if line and "-- " not in line: 122 | return MagicArgs(bind_variable=line, kwargs=kwargs) 123 | if line.endswith("--"): 124 | return MagicArgs(bind_variable=line.rstrip("--"), kwargs=kwargs) 125 | 126 | if not line: 127 | bind_variable = "_" 128 | raw = line 129 | elif line.startswith("-- "): 130 | bind_variable = "_" 131 | raw = line[3:] 132 | else: 133 | bind_variable, raw = line.split(" -- ", maxsplit=1) 134 | 135 | for arg, value in self.parse_args(shlex.split(raw)).__dict__.items(): 136 | kwargs[str(arg)] = _validate_arg_value(arg, value) 137 | 138 | return MagicArgs(bind_variable, kwargs) 139 | 140 | 141 | def _usage_string( 142 | mtyp: MagicType, args: Iterable[Text], kwargs: Iterable[Text]) -> Text: 143 | arguments = ( 144 | " ".join(f"[--{key} {key.upper()}]" for key in kwargs) + " " + 145 | " ".join(args)) 146 | 147 | if mtyp == MagicType.LINE: 148 | return f"\n```%%%(prog)s [bind_variable] -- [-h] {arguments}```" 149 | else: 150 | return f"\n```\n%%%%%(prog)s [bind_variable] -- [-h] {arguments}\ncell\n```" 151 | 152 | 153 | @dataclass(frozen=True) 154 | class _MagicSpec: 155 | """Magic function specification for argument parsing purposes.""" 156 | name: Text 157 | docstring: Text 158 | typ: MagicType 159 | 160 | args_with_no_defaults: List[Text] = field(default_factory=list) 161 | args_with_defaults: Dict[Text, Text] = field(default_factory=dict) 162 | args_descriptions: Dict[Text, Text] = field(default_factory=dict) 163 | args_types: Dict[Text, _MagicArgValueType] = field(default_factory=dict) 164 | 165 | def __post_init__(self): 166 | if (self.args_with_no_defaults or 167 | self.args_with_defaults) and not self.args_descriptions: 168 | raise MagicParsingError( 169 | "Magics have to have docstring section describing their arguments.") 170 | # pylint: disable=unsupported-membership-test 171 | for arg, typ in self.args_types.items(): 172 | if typ == bool and arg in self.args_with_no_defaults: 173 | raise MagicParsingError( 174 | "Arguments of type bool have to have a default value specified.") 175 | for arg in list(self.args_with_no_defaults) + list(self.args_with_defaults): 176 | if arg not in self.args_descriptions: 177 | raise MagicParsingError( 178 | "Magics have to have docstring section describing all of their " 179 | f"arguments; docstring missing for `{arg}`") 180 | if self.typ == MagicType.CELL and "cell" not in self.args_with_no_defaults: 181 | raise MagicParsingError( 182 | "Cell magics have to have positional argument called `cell`") 183 | # pylint: enable=unsupported-membership-test 184 | 185 | @classmethod 186 | def from_function( 187 | cls, 188 | typ: MagicType, 189 | func: Callable[..., Any], 190 | name: Optional[Text] = None) -> "_MagicSpec": 191 | """Creates _MagicSpec from compatible function.""" 192 | name = name if name else func.__name__ 193 | 194 | if not name.isidentifier(): 195 | raise MagicParsingError( 196 | f"\"{name}\" isn't valid Python identifier, see: " 197 | "https://docs.python.org/3/reference/" 198 | "lexical_analysis.html#identifiers") 199 | if not func.__doc__: 200 | raise MagicParsingError("Magics have to have docstring.") 201 | 202 | spec = getfullargspec(func) 203 | if spec.varargs or spec.varkw: 204 | raise MagicParsingError( 205 | "Magics can't have explicit variadic arguments, " 206 | "i.e. `*args` or `**kwargs`") 207 | if spec.kwonlyargs and not spec.kwonlydefaults: 208 | raise MagicParsingError( 209 | "Magics can't have keyword-only arguments without default value.") 210 | 211 | args: List[Text] = spec.args 212 | args_with_defaults: Dict[Text, Text] = {} 213 | args_types: Dict[Text, _MagicArgValueType] = {} 214 | 215 | if spec.annotations: 216 | for arg, typ_ in spec.annotations.items(): 217 | if arg == "return" or typ_ in _MagicArgValues: 218 | args_types[arg] = typ_ 219 | else: 220 | raise MagicParsingError( 221 | f"Magics can only have arguments of type {MagicArgValue}; " 222 | f"got {arg}: {typ_}") 223 | 224 | if spec.defaults: 225 | for default in reversed(spec.defaults): 226 | args_with_defaults[args.pop()] = default 227 | 228 | if spec.kwonlydefaults: 229 | for arg, default in spec.kwonlydefaults.items(): 230 | args_with_defaults[arg] = default 231 | 232 | docstring = func.__doc__ if func.__doc__ else "" 233 | 234 | args_descriptions: Dict[Text, Text] = { 235 | param.arg_name: param.description # type: ignore 236 | for param in parse(docstring).params 237 | } 238 | 239 | return cls( 240 | name, docstring, typ, args, args_with_defaults, args_descriptions, 241 | args_types) 242 | 243 | def to_parser(self) -> MagicArgumentParser: 244 | """Create an argument parser out of _MagicSpec.""" 245 | desc, *_ = self.docstring.split('\n') 246 | 247 | visible_args = self.args_with_no_defaults 248 | if self.typ == MagicType.CELL: 249 | visible_args = [a for a in self.args_with_no_defaults if a != "cell"] 250 | parser = MagicArgumentParser( 251 | prog=self.name, 252 | description=desc, 253 | usage=_usage_string( 254 | self.typ, visible_args, self.args_with_defaults.keys())) # pylint: disable=no-member 255 | 256 | for arg in self.args_with_no_defaults: # pylint: disable=not-an-iterable 257 | if self.typ == MagicType.CELL and arg == "cell": 258 | continue 259 | typ = self.args_types.get(arg, str) 260 | parser.add_argument( 261 | arg, 262 | type=typ, # type: ignore 263 | action="store", 264 | help=self.args_descriptions[arg]) 265 | for arg, default in self.args_with_defaults.items(): # pylint: disable=no-member 266 | typ = self.args_types.get(arg, str) 267 | if typ == bool: 268 | parser.add_argument( 269 | f"--{arg}", 270 | dest=arg, 271 | action="store_false" if default else "store_true", 272 | help=self.args_descriptions[arg], 273 | default=default) 274 | else: 275 | parser.add_argument( 276 | f"--{arg}", 277 | dest=arg, 278 | type=typ, # type: ignore 279 | action="store", 280 | help=self.args_descriptions[arg], 281 | default=default) 282 | return parser 283 | 284 | 285 | @dataclass(frozen=True) 286 | class Magic: 287 | """Wrapper for IPython magic.""" 288 | _spec: _MagicSpec 289 | _parser: MagicArgumentParser 290 | func: Callable[..., Any] 291 | __doc__: Text 292 | 293 | @classmethod 294 | def wrap( 295 | cls, 296 | typ: MagicType, 297 | func: Callable[..., Any], 298 | name: Optional[Text] = None) -> "Magic": 299 | """Wrap wraps a function to make it an IPython magic.""" 300 | spec = _MagicSpec.from_function(typ, func, name=name) 301 | return cls(spec, spec.to_parser(), func, func.__doc__ or cls.__doc__) 302 | 303 | def __call__(self, line: Text, cell: Optional[Text] = None) -> Any: 304 | try: 305 | args = self._parser.parse_magic_args(line, cell) 306 | res = self.func(**args.kwargs) 307 | ipython_bind_global(args.bind_variable, res) 308 | return res 309 | except ArgParserNonZeroStatus: 310 | return None 311 | -------------------------------------------------------------------------------- /picatrix/lib/magic_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # type: ignore 16 | # pylint: disable=W0212 17 | # pylint: disable=unused-argument 18 | """Test for Picatrix integration with IPython magics.""" 19 | 20 | from typing import Text 21 | 22 | import pytest 23 | 24 | from .magic import ( 25 | MagicArgParsingError, 26 | MagicArgs, 27 | MagicParsingError, 28 | MagicType, 29 | _MagicSpec, 30 | ) 31 | 32 | 33 | # Test functions 34 | def magic( 35 | a: Text, 36 | b: Text, 37 | ): 38 | """Example function. 39 | 40 | Args: 41 | a: first argument 42 | b: second argument 43 | """ 44 | 45 | 46 | def test_magicspec_valid(): 47 | """Test if _MagicSpec parses simple function.""" 48 | 49 | want = _MagicSpec( 50 | name="magic", 51 | docstring=magic.__doc__, 52 | typ=MagicType.LINE, 53 | args_with_no_defaults=["a", "b"], 54 | args_descriptions={ 55 | "a": "first argument", 56 | "b": "second argument", 57 | }, 58 | args_types={ 59 | "a": str, 60 | "b": str, 61 | }) 62 | got = _MagicSpec.from_function(MagicType.LINE, magic) 63 | 64 | assert want == got 65 | 66 | 67 | def magic_with_defaults( 68 | a: Text, 69 | b: Text = "dog", 70 | c: Text = "cat", 71 | ): 72 | """Example function. 73 | 74 | Args: 75 | a: first argument 76 | b: second argument 77 | c: third argument 78 | """ 79 | 80 | 81 | def test_magicspec_valid_with_defaults(): 82 | """Test if _MagicSpec parses simple function with default value.""" 83 | 84 | want = _MagicSpec( 85 | name="magic_with_defaults", 86 | docstring=magic_with_defaults.__doc__, 87 | typ=MagicType.LINE, 88 | args_with_no_defaults=["a"], 89 | args_with_defaults={ 90 | "b": "dog", 91 | "c": "cat", 92 | }, 93 | args_descriptions={ 94 | "a": "first argument", 95 | "b": "second argument", 96 | "c": "third argument", 97 | }, 98 | args_types={ 99 | "a": str, 100 | "b": str, 101 | "c": str, 102 | }) 103 | got = _MagicSpec.from_function(MagicType.LINE, magic_with_defaults) 104 | 105 | assert want == got 106 | 107 | 108 | # pylint: disable=missing-function-docstring 109 | def magic_no_docstring( 110 | a: Text, 111 | b: Text = "dog", 112 | ): 113 | pass 114 | 115 | 116 | # pylint: enable=missing-function-docstring 117 | 118 | 119 | def test_magicspec_no_docstring(): 120 | """Test if _MagicSpec errors on function with no docstring.""" 121 | 122 | with pytest.raises(MagicParsingError): 123 | _MagicSpec.from_function(MagicType.LINE, magic_no_docstring) 124 | 125 | 126 | def magic_bad_docstring( 127 | a: Text, 128 | b: Text = "dog", 129 | ): 130 | """Example function. 131 | 132 | Args: 133 | a: first argument 134 | """ 135 | 136 | 137 | def test_magicspec_bad_docstring(): 138 | """Test if _MagicSpec errors on function with no docstring for all args.""" 139 | 140 | with pytest.raises(MagicParsingError): 141 | _MagicSpec.from_function(MagicType.LINE, magic_bad_docstring) 142 | 143 | 144 | def magic_notstring( 145 | a: int, 146 | b: Text = "dog", 147 | c: float = 1.11, 148 | ): 149 | """Example function. 150 | 151 | Args: 152 | a: first argument 153 | b: second argument 154 | c: third argument 155 | """ 156 | 157 | 158 | def test_magicspec_nonstring(): 159 | """Test if _MagicSpec errors on non-string parameter.""" 160 | 161 | want = _MagicSpec( 162 | name="magic_notstring", 163 | docstring=magic_notstring.__doc__, 164 | typ=MagicType.LINE, 165 | args_with_no_defaults=["a"], 166 | args_with_defaults={ 167 | "b": "dog", 168 | "c": 1.11, 169 | }, 170 | args_descriptions={ 171 | "a": "first argument", 172 | "b": "second argument", 173 | "c": "third argument", 174 | }, 175 | args_types={ 176 | "a": int, 177 | "b": str, 178 | "c": float, 179 | }) 180 | got = _MagicSpec.from_function(MagicType.LINE, magic_notstring) 181 | 182 | assert want == got 183 | 184 | 185 | def magic_variadic( 186 | a: Text, 187 | *args: Text, 188 | ) -> Text: 189 | """Example function. 190 | 191 | Args: 192 | a: first argument 193 | *args: all other arguments 194 | """ 195 | 196 | 197 | def test_magicspec_variadic(): 198 | """Test if _MagicSpec errors on functions with variadic params.""" 199 | 200 | with pytest.raises(MagicParsingError): 201 | _MagicSpec.from_function(MagicType.LINE, magic_variadic) 202 | 203 | 204 | def magic_kwvariadic( 205 | a: Text, 206 | **kwargs: Text, 207 | ) -> Text: 208 | """Example function. 209 | 210 | Args: 211 | a: first argument 212 | **kwargs: all other arguments 213 | """ 214 | 215 | 216 | def test_magicspec_kwvariadic(): 217 | """Test if _MagicSpec errors on functions with keyword variadic params.""" 218 | 219 | with pytest.raises(MagicParsingError): 220 | _MagicSpec.from_function(MagicType.LINE, magic_kwvariadic) 221 | 222 | 223 | def magic_with_all_defaults( 224 | a: Text = "dog", 225 | b: Text = "cat", 226 | ): 227 | """Example function. 228 | 229 | Args: 230 | a: first argument 231 | b: second argument 232 | """ 233 | 234 | 235 | def test_magicsargsparser_empty(): 236 | """Test parsing of empty string.""" 237 | 238 | spec = _MagicSpec.from_function(MagicType.LINE, magic_with_all_defaults) 239 | parser = spec.to_parser() 240 | 241 | line = "" 242 | want = MagicArgs( 243 | bind_variable="_", kwargs={ 244 | "a": "dog", 245 | "b": "cat", 246 | }) 247 | got = parser.parse_magic_args(line, None) 248 | assert want == got 249 | 250 | 251 | def test_magicsargsparser_positional(): 252 | """Test parsing of a signal positional argument.""" 253 | spec = _MagicSpec.from_function(MagicType.LINE, magic_with_defaults) 254 | parser = spec.to_parser() 255 | 256 | line = "v -- horse" 257 | want = MagicArgs( 258 | bind_variable="v", kwargs={ 259 | "a": "horse", 260 | "b": "dog", 261 | "c": "cat", 262 | }) 263 | got = parser.parse_magic_args(line, None) 264 | assert want == got 265 | 266 | 267 | def test_magicsargsparser_mixed(): 268 | """Test parsing of a mix of positional and keyword arguments.""" 269 | spec = _MagicSpec.from_function(MagicType.LINE, magic_with_defaults) 270 | parser = spec.to_parser() 271 | 272 | line = "-- --b=boar horse" 273 | want = MagicArgs( 274 | bind_variable="_", kwargs={ 275 | "a": "horse", 276 | "b": "boar", 277 | "c": "cat", 278 | }) 279 | got = parser.parse_magic_args(line, None) 280 | assert want == got 281 | 282 | 283 | def test_magicsargsparser_notstring(): 284 | """Test parsing of non-string arguments.""" 285 | spec = _MagicSpec.from_function(MagicType.LINE, magic_notstring) 286 | parser = spec.to_parser() 287 | 288 | line = "-- --b=boar --c=2.22 10" 289 | want = MagicArgs( 290 | bind_variable="_", kwargs={ 291 | "a": 10, 292 | "b": "boar", 293 | "c": 2.22, 294 | }) 295 | got = parser.parse_magic_args(line, None) 296 | assert want == got 297 | 298 | 299 | def magic_bool_true(a: bool = True): 300 | """Example function. 301 | 302 | Args: 303 | a: first argument 304 | """ 305 | 306 | 307 | def test_magicsargsparser_bool_true(): 308 | """Test parsing of boolean argument with default equal to true.""" 309 | spec = _MagicSpec.from_function(MagicType.LINE, magic_bool_true) 310 | parser = spec.to_parser() 311 | 312 | line = "-- --a" 313 | want = MagicArgs(bind_variable="_", kwargs={"a": False}) 314 | got = parser.parse_magic_args(line, None) 315 | assert want == got 316 | 317 | line = "" 318 | want = MagicArgs(bind_variable="_", kwargs={"a": True}) 319 | got = parser.parse_magic_args(line, None) 320 | assert want == got 321 | 322 | 323 | def magic_bool_false(a: bool = False): 324 | """Example function. 325 | 326 | Args: 327 | a: first argument 328 | """ 329 | 330 | 331 | def test_magicsargsparser_bool_false(): 332 | """Test parsing of boolean argument with default equal to false.""" 333 | spec = _MagicSpec.from_function(MagicType.LINE, magic_bool_false) 334 | parser = spec.to_parser() 335 | 336 | line = "-- --a" 337 | want = MagicArgs(bind_variable="_", kwargs={"a": True}) 338 | got = parser.parse_magic_args(line, None) 339 | assert want == got 340 | 341 | line = "" 342 | want = MagicArgs(bind_variable="_", kwargs={"a": False}) 343 | got = parser.parse_magic_args(line, None) 344 | assert want == got 345 | 346 | 347 | @pytest.mark.parametrize("line", ["--", " -- ", "-- --b=cat", "11 -- cat"]) 348 | def test_magicsargsparser_error(line): 349 | """Test parsing behavior for errornous inputs.""" 350 | spec = _MagicSpec.from_function(MagicType.LINE, magic_with_defaults) 351 | parser = spec.to_parser() 352 | 353 | with pytest.raises(MagicArgParsingError): 354 | _ = parser.parse_magic_args(line, None) 355 | -------------------------------------------------------------------------------- /picatrix/lib/manager.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Class that defines the manager for all magics.""" 15 | 16 | import functools 17 | from dataclasses import dataclass 18 | from typing import Any, Callable, Dict, List, Optional, Text, Tuple, Union 19 | 20 | import pandas 21 | from IPython import get_ipython 22 | 23 | from picatrix.lib import state, utils 24 | 25 | 26 | @dataclass 27 | class Helper: 28 | """Small structure for a helper.""" 29 | function: Callable[..., Any] 30 | help: Text 31 | types: Dict[Text, Any] 32 | 33 | 34 | class MagicManager: 35 | """Manager class for Picatrix magics.""" 36 | 37 | MAGICS_DF_COLUMNS = ['name', 'description', 'line', 'cell', 'function'] 38 | 39 | _magics: Dict[Text, Callable[[Text, Text], Text]] = {} 40 | _helpers: Dict[Text, Helper] = {} 41 | 42 | @classmethod 43 | def clear_helpers(cls): 44 | """Clear all helper registration.""" 45 | for helper_name in cls._helpers: 46 | try: 47 | utils.ipython_remove_global(helper_name) 48 | except KeyError: 49 | pass 50 | cls._helpers = {} 51 | 52 | @classmethod 53 | def clear_magics(cls): 54 | """Clear all magic registration.""" 55 | magics = list(cls._magics.keys()) 56 | for magic_name in magics: 57 | cls.deregister_magic(magic_name) 58 | 59 | @classmethod 60 | def deregister_helper(cls, helper_name: Text): 61 | """Remove a helper from the registration. 62 | 63 | Args: 64 | helper_name (str): the name of the helper to remove. 65 | 66 | Raises: 67 | KeyError: if the helper is not registered. 68 | """ 69 | if helper_name not in cls._helpers: 70 | raise KeyError(f'Helper [{helper_name}] is not registered.') 71 | 72 | _ = cls._helpers.pop(helper_name) 73 | try: 74 | utils.ipython_remove_global(helper_name) 75 | except KeyError: 76 | pass 77 | 78 | @classmethod 79 | def deregister_magic(cls, magic_name: Text): 80 | """Removes a magic from the registration. 81 | 82 | Args: 83 | magic_name (str): the name of the magic to remove. 84 | 85 | Raises: 86 | KeyError: if the magic is not registered. 87 | """ 88 | if magic_name not in cls._magics: 89 | raise KeyError(f'Magic [{magic_name}] is not registered.') 90 | 91 | _ = cls._magics.pop(magic_name) 92 | try: 93 | utils.ipython_remove_global(f'{magic_name}_func') 94 | except KeyError: 95 | pass 96 | 97 | # Attempt to remove the magic definition. 98 | ip = get_ipython() 99 | magics_manager = ip.magics_manager 100 | 101 | if not hasattr(magics_manager, 'magics'): 102 | return 103 | 104 | line_magics = magics_manager.magics.get('line', {}) 105 | if magic_name in line_magics: 106 | _ = magics_manager.magics.get('line').pop(magic_name) 107 | 108 | cell_magics = magics_manager.magics.get('cell', {}) 109 | if magic_name in cell_magics: 110 | _ = magics_manager.magics.get('cell').pop(magic_name) 111 | 112 | @classmethod 113 | def get_helper(cls, helper_name: Text) -> Optional[Callable[..., Any]]: 114 | """Return a helper function from the registration.""" 115 | return cls._magics.get(helper_name) 116 | 117 | @classmethod 118 | def get_magic(cls, magic_name: Text) -> Callable[[Text, Text], Text]: 119 | """Return a magic function from the registration.""" 120 | return cls._magics.get(magic_name) 121 | 122 | @classmethod 123 | def get_helper_info( 124 | cls, 125 | as_pandas: Optional[bool] = True 126 | ) -> Union[pandas.DataFrame, List[Tuple[Text, Text]]]: 127 | """Get a list of all the registered helpers. 128 | 129 | Args: 130 | as_pandas (bool): boolean to determine whether to receive the results 131 | as a list of tuple or a pandas DataFrame. Defaults to True. 132 | 133 | Returns: 134 | Either a pandas DataFrame or a list of tuples, depending on the 135 | as_pandas boolean. 136 | """ 137 | if not as_pandas: 138 | return [(name, helper.help) for name, helper in cls._helpers.items()] 139 | 140 | lines = [] 141 | for name, helper in cls._helpers.items(): 142 | hints = helper.types 143 | hint_strings = [] 144 | for key, value in hints.items(): 145 | value_string = getattr(value, '__name__', str(value)) 146 | hint_strings.append(f'{key} [{value_string}]') 147 | helper_string = ', '.join(hint_strings) 148 | 149 | lines.append( 150 | { 151 | 'name': name, 152 | 'help': helper.help, 153 | 'arguments': helper_string, 154 | }) 155 | return pandas.DataFrame(lines) 156 | 157 | @classmethod 158 | def get_magic_info( 159 | cls, 160 | as_pandas: Optional[bool] = True 161 | ) -> Union[pandas.DataFrame, List[Tuple[Text, Text]]]: 162 | """Get a list of all magics. 163 | 164 | Args: 165 | as_pandas (bool): boolean to determine whether to receive the results 166 | as a list of tuples or a pandas DataFrame. Defaults to True. 167 | 168 | Returns: 169 | Either a pandas DataFrame or a list of tuples, depending on the as_pandas 170 | boolean. 171 | """ 172 | if not as_pandas: 173 | return [ 174 | (x.magic_name, x.__doc__.split('\n')[0]) 175 | for x in iter(cls._magics.values()) 176 | ] 177 | 178 | entries = [] 179 | for magic_name, magic_class in iter(cls._magics.items()): 180 | description = magic_class.__doc__.split('\n')[0] 181 | magic_dict = { 182 | 'name': magic_name, 183 | 'cell': f'%%{magic_name}', 184 | 'line': f'%{magic_name}', 185 | 'function': f'{magic_name}_func', 186 | 'description': description 187 | } 188 | entries.append(magic_dict) 189 | 190 | df = pandas.DataFrame(entries) 191 | return df[cls.MAGICS_DF_COLUMNS].sort_values('name') 192 | 193 | @classmethod 194 | def register_helper( 195 | cls, name: Text, helper: Any, typing_help: Dict[Text, Any]): 196 | """Register a picatrix helper function. 197 | 198 | Args: 199 | name (str): the name of the helper function. 200 | helper (function): the helper function to register. 201 | typing_help (dict): dict with the arguments and their types. 202 | 203 | Raises: 204 | KeyError: if the helper is already registered. 205 | """ 206 | if name in cls._helpers: 207 | raise KeyError(f'The helper [{name}] is already registered.') 208 | doc_string = helper.__doc__ 209 | if doc_string: 210 | help_string = doc_string.split('\n')[0] 211 | else: 212 | help_string = 'No help string supplied.' 213 | 214 | cls._helpers[name] = Helper( 215 | function=helper, help=help_string, types=typing_help) 216 | 217 | @classmethod 218 | def register_magic( 219 | cls, 220 | function: Callable[[Text, Text], Text], 221 | conditional: Callable[[], bool] = None): 222 | """Register magic function as a magic in picatrix. 223 | 224 | Args: 225 | function (function): the function to register as a line and a 226 | cell magic. 227 | conditional (function): a function that should return a bool, used to 228 | determine whether to register magic or not. This can be used by 229 | magics to determine whether a magic should be registered or not, for 230 | instance basing that on whether the notebook is able to reach the 231 | required service, or whether a connection to a client can be achieved, 232 | etc. This is optional and if not provided a magic will be registered. 233 | 234 | Raises: 235 | KeyError: if the magic is already registered. 236 | """ 237 | if conditional and not conditional(): 238 | return 239 | 240 | magic_name = function.magic_name 241 | 242 | if magic_name in cls._magics: 243 | raise KeyError(f'The magic [{magic_name}] is already registered.') 244 | 245 | ip = get_ipython() 246 | if ip: 247 | ip.register_magic_function( 248 | function, magic_kind='line_cell', magic_name=magic_name) 249 | 250 | cls._magics[magic_name] = function 251 | function_name = f'{magic_name}_func' 252 | 253 | def capture_output(function, name): 254 | """A function that wraps around magic functions to capture output.""" 255 | 256 | @functools.wraps(function) 257 | def wrapper(*args, **kwargs): 258 | function_output = function(*args, **kwargs) 259 | state_obj = state.state() 260 | return state_obj.set_output(function_output, magic_name=name) 261 | 262 | return wrapper 263 | 264 | _ = utils.ipython_bind_global( 265 | function_name, capture_output(function.fn, function_name)) 266 | -------------------------------------------------------------------------------- /picatrix/lib/manager_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Copyright 2020 Google LLC 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # https://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """Tests for the pixatrix manager.""" 16 | import typing 17 | 18 | import mock 19 | import pytest 20 | 21 | from picatrix.lib import manager, utils 22 | 23 | manager.get_ipython = mock.MagicMock() 24 | utils.get_ipython = mock.MagicMock() 25 | 26 | 27 | def test_registration(): 28 | """Test registering a magic and getting a copy of it and de-registering.""" 29 | manager.MagicManager.clear_magics() 30 | 31 | def my_magic(cell=None, line=None): 32 | """This is a magic.""" 33 | if not cell: 34 | cell = 'foo' 35 | if not line: 36 | line = 'bar' 37 | return f'{cell}{line}' 38 | 39 | my_magic.magic_name = 'magical_function' 40 | my_magic.fn = my_magic 41 | manager.MagicManager.register_magic(my_magic) 42 | 43 | magic_from_manager = manager.MagicManager.get_magic('magical_function') 44 | assert magic_from_manager() == 'foobar' 45 | 46 | my_magic.magic_name = 'other_magic' 47 | 48 | def conditional(): 49 | return False 50 | 51 | manager.MagicManager.register_magic(my_magic, conditional=conditional) 52 | magic_from_manager = manager.MagicManager.get_magic('other_magic') 53 | assert magic_from_manager is None 54 | 55 | manager.MagicManager.register_magic(my_magic) 56 | magic_from_manager = manager.MagicManager.get_magic('other_magic') 57 | assert magic_from_manager() == 'foobar' 58 | 59 | manager.MagicManager.deregister_magic('other_magic') 60 | magic_from_manager = manager.MagicManager.get_magic('other_magic') 61 | assert magic_from_manager is None 62 | 63 | manager.MagicManager.deregister_magic('magical_function') 64 | magic_from_manager = manager.MagicManager.get_magic('magical_function') 65 | assert magic_from_manager is None 66 | 67 | with pytest.raises(KeyError): 68 | manager.MagicManager.deregister_magic('does_not_exist') 69 | 70 | 71 | def test_magic_info(): 72 | """Test the get_magic_info.""" 73 | # Start by clearing the current registration. 74 | manager.MagicManager.clear_magics() 75 | 76 | def magical_func(): 77 | """This is a magical function that returns pure magic.""" 78 | return 'magic' 79 | 80 | magical_func.magic_name = 'magical_function' 81 | magical_func.fn = magical_func 82 | manager.MagicManager.register_magic(magical_func) 83 | 84 | def second_magic(): 85 | """This is even more magic.""" 86 | return 'fab' 87 | 88 | second_magic.magic_name = 'some_magic' 89 | second_magic.fn = second_magic 90 | manager.MagicManager.register_magic(second_magic) 91 | 92 | def other_magic(): 93 | """Could this be it?""" 94 | return 'true magic' 95 | 96 | other_magic.magic_name = 'other_magic' 97 | other_magic.fn = other_magic 98 | manager.MagicManager.register_magic(other_magic) 99 | 100 | info_df = manager.MagicManager.get_magic_info(as_pandas=True) 101 | assert len(info_df) == 3 102 | assert not info_df[info_df.name == 'other_magic'].empty 103 | 104 | desc_set = set(info_df.description.unique()) 105 | expected_set = set( 106 | [ 107 | 'Could this be it?', 'This is even more magic.', 108 | 'This is a magical function that returns pure magic.' 109 | ]) 110 | 111 | assert desc_set == expected_set 112 | 113 | entries = manager.MagicManager.get_magic_info(as_pandas=False) 114 | assert len(entries) == 3 115 | names = [x[0] for x in entries] 116 | assert 'other_magic' in names 117 | assert 'some_magic' in names 118 | 119 | 120 | def test_helper_registration(): 121 | """Test registering helpers.""" 122 | 123 | def helper_func(stuff: typing.Text): 124 | """This is a helper function that helps.""" 125 | back = f'{stuff}' * 3 126 | return f'{back}\n\nwasn\'t this helpful?' 127 | 128 | helper_dict = typing.get_type_hints(helper_func) 129 | manager.MagicManager.register_helper( 130 | name='helper_func', helper=helper_func, typing_help=helper_dict) 131 | 132 | df = manager.MagicManager.get_helper_info() 133 | assert not df[df.name == 'helper_func'].empty 134 | series = df.set_index('name').loc['helper_func'] 135 | assert series.arguments == 'stuff [str]' 136 | -------------------------------------------------------------------------------- /picatrix/lib/namespace.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Types and functions defining Picatrix namespacing.""" 15 | 16 | from difflib import get_close_matches 17 | from inspect import cleandoc 18 | from types import SimpleNamespace 19 | from typing import ( 20 | Any, 21 | Callable, 22 | Dict, 23 | Generic, 24 | Iterable, 25 | Iterator, 26 | List, 27 | Optional, 28 | Text, 29 | Tuple, 30 | TypeVar, 31 | Union, 32 | ) 33 | 34 | import pandas 35 | 36 | from .error import Error 37 | from .ipython import add_cell_magic, add_line_magic, delete_magic 38 | from .magic import Magic, MagicType 39 | 40 | 41 | class NamespaceKeyError(Error, KeyError): 42 | """Raised when Namespace attempted to operate on invalid key.""" 43 | 44 | 45 | class NamespaceKeyExistsError(NamespaceKeyError): 46 | """Raised when add operation was called on already existing key.""" 47 | 48 | 49 | class NamespaceKeyMissingError(NamespaceKeyError): 50 | """Raised when non-existing key is attempted to be accessed in a Namespace.""" 51 | 52 | def __init__(self, key: Text, other_keys: Iterable[Text]): 53 | matches = get_close_matches(key, other_keys, 1) 54 | if not matches: 55 | super().__init__(f"{key} does not exist") 56 | else: 57 | super().__init__( 58 | f"{key} does not exist; " 59 | f"did you mean \"{matches[0]}\"") 60 | 61 | 62 | A = TypeVar("A") # pylint: disable=invalid-name 63 | 64 | 65 | def _join_key(*args: Text): 66 | return ".".join(args) 67 | 68 | 69 | def _as_df_record(name: Text, item: Any, with_doc: bool, 70 | with_values: bool) -> Dict[Text, Text]: 71 | """Create a record compatible with pandas.DataFrame.from_records func. 72 | 73 | Args: 74 | name: human readable name of the item 75 | item: item to be translated into the row 76 | with_doc: if True, add whole docstring as a column; 1st line only otherwise 77 | 78 | Returns: 79 | Dict[Text, Text]: dictionary compatible with pandas.DataFrame.from_records 80 | """ 81 | typ = "Namespace" if isinstance(item, Namespace) else str(type(item)) 82 | doc = str(item.__doc__) if item.__doc__ else "" 83 | desc, *_ = doc.split("\n") 84 | 85 | record = { 86 | "Name": name, 87 | "Type": typ, 88 | "Description": desc, 89 | } 90 | 91 | if with_doc: 92 | record["Docstring"] = doc 93 | if with_values: 94 | record["Value"] = str(item) # type: ignore 95 | 96 | return record 97 | 98 | 99 | class Namespace(SimpleNamespace, Generic[A]): 100 | """Key-value type of structure with items accessible as attributes.""" 101 | 102 | name: Text 103 | __dict__: Dict[Text, A] 104 | 105 | def __init__(self, name: Text, docstring: Text, **kwargs: A): 106 | super().__init__(**kwargs) 107 | self.__doc__ = cleandoc(docstring) 108 | self.name = name 109 | 110 | def __setattr__(self, key: Text, value: A): 111 | if not key.isidentifier(): 112 | raise NamespaceKeyError( 113 | f"\"{key}\" isn't valid Python identifier, see: " 114 | "https://docs.python.org/3/reference/" 115 | "lexical_analysis.html#identifiers") 116 | elif key in self: 117 | raise NamespaceKeyExistsError( 118 | f"\"{key}\" already exists; remove it " 119 | "first with `del` operator, " 120 | f"e.g. `del {self.name}.{key}`") 121 | else: 122 | super().__setattr__(key, value) 123 | 124 | def __getattr__(self, key: Text) -> A: 125 | if key in self: 126 | return super().__getattr__(key) 127 | else: 128 | raise NamespaceKeyMissingError(key, self.keys()) 129 | 130 | def __delattr__(self, key: Text): 131 | if key in self: 132 | return super().__delattr__(key) 133 | else: 134 | raise NamespaceKeyMissingError(key, self.keys()) 135 | 136 | def __iter__(self) -> Iterator[Text]: 137 | return iter(self.__dict__) 138 | 139 | def __contains__(self, key: Text): 140 | return self.__dict__.__contains__(key) 141 | 142 | def keys(self) -> Iterator[Text]: 143 | """Iterator over all of the keys in the namespace.""" 144 | return iter(self.__dict__.keys()) 145 | 146 | def values(self) -> Iterator[A]: 147 | """Iterator over all of the values in the namespace.""" 148 | return iter(self.__dict__.values()) 149 | 150 | def items(self) -> Iterator[Tuple[Text, A]]: 151 | """Iterator over all of the key-value pairs in the namespace.""" 152 | for key, value in self.__dict__.items(): 153 | yield (key, value) 154 | 155 | def as_records(self, 156 | with_doc: bool = False, 157 | with_values: bool = False) -> List[Dict[Text, Text]]: 158 | """Make into records to be used with pandas.DataFrame.from_records func.""" 159 | records = [_as_df_record(self.name, self, with_doc, False)] 160 | for key, value in self.items(): 161 | if key == "name" or key.startswith("_"): 162 | continue 163 | 164 | if isinstance(value, Namespace): 165 | records.extend(value.as_records(with_doc, with_values)) 166 | else: 167 | records.append( 168 | _as_df_record( 169 | _join_key(self.name, key), value, with_doc, with_values)) 170 | return records 171 | 172 | def to_frame( 173 | self, 174 | with_doc: bool = False, 175 | with_values: bool = False) -> pandas.DataFrame: 176 | """Make namespace into pandas.DataFrame.""" 177 | return pandas.DataFrame.from_records( # type: ignore 178 | self.as_records(with_doc, with_values)) 179 | 180 | def search(self, keyword: Text) -> pandas.DataFrame: 181 | """Search namespace for elements containing keyword.""" 182 | df = self.to_frame(with_doc=True) 183 | return df[df.Name.str.contains(keyword) | # type: ignore 184 | df.Docstring.str.contains(keyword)] # type: ignore 185 | 186 | def _add(self, key: Text, value: A): 187 | """Adds a new value under the key. 188 | 189 | Raises: 190 | NamespaceKeyExistsError: when required key already exists 191 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 192 | """ 193 | setattr(self, key, value) 194 | 195 | def get(self, key: Text, default: A) -> A: 196 | """Return the value for key if key is in the namespace, else default.""" 197 | if key in self: 198 | return getattr(self, key) 199 | else: 200 | return default 201 | 202 | def delete(self, key: Text): 203 | """Deletes a key from the namespace.""" 204 | self.__delattr__(key) 205 | 206 | 207 | Function = Callable[..., Any] 208 | """Type representation of a function supported by Picatrix.""" 209 | 210 | 211 | class MagicNamespace(Namespace[Function]): 212 | """Namespace hosting IPython magics.""" 213 | 214 | def __delattr__(self, key: Text): 215 | super().__delattr__(key) 216 | delete_magic(key) 217 | 218 | def add_line_magic(self, func: Function, name: Optional[Text] = None): 219 | """Adds new line magic (`%magic`) to the namespace. 220 | 221 | Args: 222 | func: function to be made into magic 223 | name: optional custom magic name, if None function name will be used 224 | 225 | Raises: 226 | NamespaceKeyExistsError: when required key already exists 227 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 228 | MagicParsingError: when provided function is an invalid magic 229 | """ 230 | name = name if name else func.__name__ 231 | 232 | magic = Magic.wrap(MagicType.LINE, func, name=name) 233 | self._add(name, magic.func) 234 | 235 | add_line_magic(name, magic) 236 | 237 | def add_cell_magic(self, func: Function, name: Optional[Text] = None): 238 | """Adds new cell magic (`%%magic`) to the namespace. 239 | 240 | Args: 241 | func: function to be made into magic 242 | name: optional custom magic name, if None function name will be used 243 | 244 | Raises: 245 | NamespaceKeyExistsError: when required key already exists 246 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 247 | MagicParsingError: when provided function is an invalid magic 248 | """ 249 | name = name if name else func.__name__ 250 | 251 | magic = Magic.wrap(MagicType.CELL, func, name=name) 252 | self._add(name, magic.func) 253 | 254 | add_cell_magic(name, magic) 255 | 256 | 257 | class FeatureNamespace(Namespace[Function]): 258 | """Namespace hosting functionalities of the specific feature.""" 259 | 260 | def add_function(self, func: Function, name: Optional[Text] = None): 261 | """Adds a new function to the namespace. 262 | 263 | Raises: 264 | NamespaceKeyExistsError: when required key already exists 265 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 266 | """ 267 | if name is None: 268 | name = func.__name__ 269 | self._add(name, func) 270 | 271 | 272 | class RootNamespace(Namespace[Union[FeatureNamespace, Function]]): 273 | """Namespace hosting all Picatrix funcionalities.""" 274 | 275 | magic: MagicNamespace 276 | 277 | def __init__( 278 | self, name: Text, docstring: Text, **kwargs: Union[FeatureNamespace, 279 | Function]): 280 | super().__init__(name, docstring, **kwargs) 281 | self.magic = MagicNamespace( 282 | _join_key(self.name, "magic"), 283 | "Namespace holding all IPython magics registered by Picatrix.") 284 | 285 | def add_function(self, func: Function, name: Optional[Text] = None): 286 | """Adds a new function to the namespace. 287 | 288 | Raises: 289 | NamespaceKeyExistsError: when required key already exists 290 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 291 | """ 292 | if name is None: 293 | name = func.__name__ 294 | self._add(name, func) 295 | 296 | def add_line_magic(self, func: Function, name: Optional[Text] = None): 297 | """Adds new line magic (`%magic`) to the namespace. 298 | 299 | Args: 300 | func: function to be made into magic 301 | name: optional custom magic name, if None function name will be used 302 | 303 | Raises: 304 | NamespaceKeyExistsError: when required key already exists 305 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 306 | MagicParsingError: when provided function is an invalid magic 307 | """ 308 | self.magic.add_line_magic(func, name) 309 | 310 | def add_cell_magic(self, func: Function, name: Optional[Text] = None): 311 | """Adds new cell magic (`%%magic`) to the namespace. 312 | 313 | Args: 314 | func: function to be made into magic 315 | name: optional custom magic name, if None function name will be used 316 | 317 | Raises: 318 | NamespaceKeyExistsError: when required key already exists 319 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 320 | MagicParsingError: when provided function is an invalid magic 321 | """ 322 | self.magic.add_cell_magic(func, name) 323 | 324 | def add_namespace( 325 | self, name: Text, docstring: Optional[Text] = None) -> FeatureNamespace: 326 | """Adds a new namespace to the root namespace. 327 | 328 | Args: 329 | name: name of the new namespace 330 | docstring: optional docstring for the new namespace 331 | 332 | Returns: 333 | new namespace 334 | 335 | Raises: 336 | NamespaceKeyExistsError: when required key already exists 337 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 338 | """ 339 | key = _join_key(self.name, name) 340 | if not docstring: 341 | docstring = f"`{key}` contains all Picatrix features of {name}" 342 | 343 | ns = FeatureNamespace(name=key, docstring=docstring) 344 | self._add(name, ns) 345 | return ns 346 | 347 | 348 | class FeatureContext(Namespace[Any]): 349 | """Namespace hosting runtime parameters of the specific feature.""" 350 | 351 | def add(self, key: Text, value: Any): 352 | """Adds a new value under the key. 353 | 354 | Raises: 355 | NamespaceKeyExistsError: when required key already exists 356 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 357 | """ 358 | self._add(key, value) 359 | 360 | 361 | class RootContext(Namespace[FeatureContext]): 362 | """Namespace hosting all Picatrix runtime parameters, i.e. context.""" 363 | 364 | def add_namespace( 365 | self, name: Text, docstring: Optional[Text] = None) -> FeatureContext: 366 | """Adds a new context to the root context. 367 | 368 | Args: 369 | name: name of the new subcontext 370 | docstring: optional docstring for the new subcontext 371 | 372 | Returns: 373 | new subcontext 374 | 375 | Raises: 376 | NamespaceKeyExistsError: when required key already exists 377 | NamespaceKeyError: when key is invalid, e.g. isn't Python identifier 378 | """ 379 | key = _join_key(self.name, name) 380 | if not docstring: 381 | docstring = f"`{key}` contains all Picatrix context for {name}" 382 | 383 | ctx = FeatureContext(name=key, docstring=docstring) 384 | self._add(name, ctx) 385 | return ctx 386 | -------------------------------------------------------------------------------- /picatrix/lib/namespace_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # type: ignore 16 | # pylint: disable=W0212 17 | """Test for Picatrix namespacing.""" 18 | 19 | from typing import Union 20 | 21 | import pandas as pd 22 | import pytest 23 | 24 | from .namespace import Namespace, NamespaceKeyError 25 | 26 | 27 | def test_invalid_key(): 28 | """Test if namespace throws an error on invalid key names.""" 29 | n = Namespace[int](name="n", docstring="Example namespace") 30 | with pytest.raises(NamespaceKeyError): 31 | n._add("1aa", 1) 32 | 33 | 34 | def test_similar_key(): 35 | """Test if namespace throws an helpful error on typos.""" 36 | n = Namespace[int](name="n", docstring="Example namespace") 37 | n._add("apple", 1) 38 | n._add("orange", 2) 39 | n._add("carrot", 3) 40 | 41 | with pytest.raises(NamespaceKeyError) as exc: 42 | _ = n.aple 43 | 44 | assert "did you mean \"apple\"" in str(exc) 45 | 46 | 47 | def test_to_frame(): 48 | """Test namespaces .to_frame functionality for nested namespaces.""" 49 | n = Namespace[Union[int, Namespace[int]]]( 50 | name="n", docstring="Example namespace") 51 | n._add("apple", 1) 52 | n._add("orange", 2) 53 | n._add("carrot", 3) 54 | n._add("nn", Namespace[int](name="n.nn", docstring="Nested namespace")) 55 | n.nn._add("dog", 11) 56 | n._add("tomato", 4) 57 | 58 | a = 1 59 | inttyp = str(type(a)) 60 | intdoc = a.__doc__ 61 | intdesc, *_ = intdoc.split("\n") 62 | 63 | want = pd.DataFrame.from_records( 64 | [ 65 | { 66 | "Name": "n", 67 | "Type": "Namespace", 68 | "Description": "Example namespace", 69 | "Docstring": "Example namespace", 70 | }, 71 | { 72 | "Name": "n.apple", 73 | "Type": inttyp, 74 | "Description": intdesc, 75 | "Docstring": intdoc, 76 | }, 77 | { 78 | "Name": "n.orange", 79 | "Type": inttyp, 80 | "Description": intdesc, 81 | "Docstring": intdoc, 82 | }, 83 | { 84 | "Name": "n.carrot", 85 | "Type": inttyp, 86 | "Description": intdesc, 87 | "Docstring": intdoc, 88 | }, 89 | { 90 | "Name": "n.nn", 91 | "Type": "Namespace", 92 | "Description": "Nested namespace", 93 | "Docstring": "Nested namespace", 94 | }, 95 | { 96 | "Name": "n.nn.dog", 97 | "Type": inttyp, 98 | "Description": intdesc, 99 | "Docstring": intdoc, 100 | }, 101 | { 102 | "Name": "n.tomato", 103 | "Type": inttyp, 104 | "Description": intdesc, 105 | "Docstring": intdoc, 106 | }, 107 | ]) 108 | got = n.to_frame(with_doc=True) 109 | assert want.equals(got) 110 | -------------------------------------------------------------------------------- /picatrix/lib/state.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Class that defines the state, which .""" 15 | import logging 16 | import threading 17 | from typing import Any, Dict, Optional, Text 18 | 19 | from picatrix.lib import utils 20 | 21 | logger = logging.getLogger('picatrix.state') 22 | 23 | __state = None 24 | _LOCK = threading.Lock() 25 | 26 | 27 | def state(refresh_state: bool = False): 28 | """Property that returns a state object.""" 29 | # pylint: disable=global-statement 30 | # Making sure we have only one state object. 31 | global __state 32 | 33 | with _LOCK: 34 | if refresh_state or __state is None: 35 | __state = State() 36 | 37 | return __state 38 | 39 | 40 | class State: 41 | """Picatrix state object.""" 42 | 43 | _cache: Dict[Text, Any] = {} 44 | _last_output: Any 45 | _last_magic: Text 46 | 47 | def __init__(self): 48 | self._last_output = None 49 | self._last_magic = '' 50 | 51 | @property 52 | def last_magic(self): 53 | """A property that returns the last magic that was run.""" 54 | return self._last_magic 55 | 56 | @property 57 | def last_output(self): 58 | """A property that returns the last output from a magic.""" 59 | return self._last_output 60 | 61 | def add_to_cache(self, name: Text, value: Any): 62 | """Add a value to the cache or update value if it exists. 63 | 64 | Args: 65 | name (str): name of the value in the cache. 66 | value (object): the value to be stored in the cache. 67 | """ 68 | with _LOCK: 69 | self._cache[name] = value 70 | 71 | def get_from_cache(self, name: Text, default: Optional[Any] = None) -> Any: 72 | """Get a value from the cache. 73 | 74 | Args: 75 | name (str): name of the value in the cache to retrieve. 76 | default (object): if the value does not exist, defines the default 77 | value to return. This is optional and returns None if not defined. 78 | 79 | Returns: 80 | The value from the cache if it exists, otherwise the default value. 81 | """ 82 | with _LOCK: 83 | return self._cache.get(name, default) 84 | 85 | def remove_from_cache(self, name: Text): 86 | """Removes a value from the cache if it exists.""" 87 | with _LOCK: 88 | if name in self._cache: 89 | del self._cache[name] 90 | 91 | def set_output( 92 | self, 93 | output: Any, 94 | magic_name: Text, 95 | bind_to: Optional[Text] = '') -> Optional[Any]: 96 | """Sets an output from a magic and stores it in the namespace if needed. 97 | 98 | Args: 99 | output (object): the output from the magic as executed. 100 | magic_name (str): the name of the magic that was used. 101 | bind_to (str): optional name of a variable. If this is provided 102 | the output is omitted and variable is stored in the namespace using 103 | the name provided here. 104 | 105 | Returns: 106 | Returns the output object. 107 | """ 108 | with _LOCK: 109 | self._last_output = output 110 | self._last_magic = magic_name 111 | 112 | if bind_to: 113 | _ = utils.ipython_bind_global(name=bind_to, value=output) 114 | return None 115 | 116 | return output 117 | -------------------------------------------------------------------------------- /picatrix/lib/state_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Copyright 2020 Google LLC 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # https://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """Tests for the pixatrix state object.""" 16 | from picatrix.lib import testlib 17 | 18 | 19 | def test_getting_last_magic(): 20 | """Test running a magic.""" 21 | ip = testlib.InteractiveTest.get_shell() 22 | ip.run_cell(raw_cell='from picatrix.magics import common') 23 | res = ip.run_cell(raw_cell='common.picatrixmagics("")') 24 | output_df = res.result 25 | #output_df = ip.run_line_magic(magic_name='picatrixmagics', line='') 26 | 27 | code = ( 28 | 'df = common.last_output("")\n' 29 | 'names = list(df.name.unique())\n' 30 | 'sorted(names)') 31 | 32 | names = list(output_df['name'].unique()) 33 | testlib.InteractiveTest.run_and_compare(code, expected_return=sorted(names)) 34 | 35 | 36 | def test_working_with_cache(): 37 | """Test working with the cache.""" 38 | ip = testlib.InteractiveTest.get_shell() 39 | res = ip.run_cell( 40 | raw_cell=('from picatrix.lib import state\n' 41 | 'state.state()\n')) 42 | assert res.success 43 | state_obj = res.result 44 | state_obj.add_to_cache('foobar', 1234) 45 | 46 | assert state_obj.get_from_cache('foobar') == 1234 47 | 48 | res = ip.run_cell( 49 | raw_cell=( 50 | 'state_obj = state.state()\n' 51 | 'state_obj.get_from_cache("foobar")\n')) 52 | assert res.success 53 | assert res.result == 1234 54 | 55 | res = ip.run_cell(raw_cell='state_obj.remove_from_cache("foobar")') 56 | assert res.success 57 | 58 | res = ip.run_cell(raw_cell='state_obj.get_from_cache("foobar")') 59 | assert res.success 60 | assert res.result is None 61 | -------------------------------------------------------------------------------- /picatrix/lib/testlib.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Test library for picatrix.""" 15 | from typing import Any, Text 16 | 17 | from IPython.terminal.interactiveshell import TerminalInteractiveShell 18 | from IPython.testing.globalipapp import start_ipython 19 | 20 | from picatrix.lib import manager 21 | 22 | manager.MagicManager.is_test = True 23 | 24 | 25 | class InteractiveTest: 26 | """Class that helps with interactive tests for picatrix.""" 27 | 28 | _shell: TerminalInteractiveShell = None 29 | 30 | @classmethod 31 | def _init_shell(cls): 32 | """Initialize the shell.""" 33 | ip = start_ipython() 34 | 35 | ip.run_cell(raw_cell='from picatrix import notebook_init') 36 | ip.run_cell(raw_cell='notebook_init.init()') 37 | cls._shell = ip 38 | 39 | @classmethod 40 | def get_shell(cls) -> TerminalInteractiveShell: 41 | """Return an interactive shell, initializing it for the first time.""" 42 | if cls._shell: 43 | return cls._shell 44 | cls._init_shell() 45 | return cls._shell 46 | 47 | @classmethod 48 | def run_and_compare(cls, code: Text, expected_return: Any): 49 | """Run code in shell, get returns and assert against expected returns.""" 50 | shell = cls.get_shell() 51 | res = shell.run_cell(raw_cell=code) 52 | assert res.success 53 | assert res.result == expected_return 54 | -------------------------------------------------------------------------------- /picatrix/lib/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Common reusable components for picatrix.""" 15 | import sys 16 | from typing import Any, Optional, Text 17 | 18 | from IPython import get_ipython 19 | from IPython.display import clear_output 20 | 21 | 22 | def ask_question(question: Text, input_type: Optional[Any] = Text) -> Any: 23 | """Asks a question and returns the answer. 24 | 25 | Args: 26 | question (str): the question to be asked. 27 | input_type (object): the type of the input data. 28 | 29 | Raises: 30 | TypeError: if the input type is not supported. 31 | 32 | Returns: 33 | object: an object of type "input_type" read back as an answer. 34 | """ 35 | print(question) 36 | answer_line = sys.stdin.readline() 37 | answer_line = answer_line.strip() 38 | 39 | if input_type == str: 40 | return answer_line 41 | 42 | if input_type in (int, float): 43 | return input_type(answer_line) 44 | 45 | raise TypeError('Only support str, int and float as input types') 46 | 47 | 48 | def clear_notebook_output(): 49 | """Clears the output cell from the notebook.""" 50 | clear_output(wait=True) 51 | 52 | 53 | def ipython_bind_global(name: Text, value: Any) -> Any: 54 | """Binds the name to a Python object denoted by value. 55 | 56 | Args: 57 | name (str): Variable name. 58 | value (object): Python object to bind to name 59 | 60 | Returns: 61 | Returns the value bound to the given name. 62 | """ 63 | ip = get_ipython() 64 | if ip and name: 65 | ip.push({name: value}) 66 | return value 67 | 68 | 69 | def ipython_get_global(name: Text) -> Any: 70 | """Returns the Python object bound to the given name in the user namespace. 71 | 72 | Args: 73 | name (str): Variable name. 74 | 75 | Returns: 76 | Returns the value bound to the given name. 77 | 78 | Raises: 79 | KeyError: if the variable is not stored in the namespace. 80 | """ 81 | ip = get_ipython() 82 | return ip.all_ns_refs[0][name] 83 | 84 | 85 | def ipython_remove_global(name: Text): 86 | """Removes a Python object that is bound to the user namespace. 87 | 88 | Args: 89 | name (str): Variable name. 90 | 91 | Raises: 92 | KeyError: if the variable is not stored in the namespace. 93 | """ 94 | ip = get_ipython() 95 | namespace = ip.all_ns_refs[0] 96 | 97 | if name not in namespace: 98 | raise KeyError(f'The variable {name} is not currently in the namespace.') 99 | 100 | _ = namespace.pop(name) 101 | -------------------------------------------------------------------------------- /picatrix/lib/utils_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Copyright 2020 Google LLC 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # https://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """Tests for the picatrix utils.""" 16 | import io 17 | 18 | from picatrix.lib import utils 19 | 20 | 21 | def test_asking_questions(monkeypatch): 22 | """Test the utils ask_question function..""" 23 | expected_answer = 'picatrix' 24 | monkeypatch.setattr('sys.stdin', io.StringIO(expected_answer)) 25 | text = utils.ask_question('What is your name') 26 | assert text == expected_answer 27 | 28 | expected_answer = 234.12 29 | monkeypatch.setattr('sys.stdin', io.StringIO('234.12')) 30 | number = utils.ask_question('What is your number', input_type=float) 31 | assert number == expected_answer 32 | -------------------------------------------------------------------------------- /picatrix/magics/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Class that loads up all magic libraries.""" 15 | 16 | from picatrix.magics import common, timesketch 17 | -------------------------------------------------------------------------------- /picatrix/magics/common.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Class that defines common picatrix magics.""" 15 | 16 | from typing import Any, Optional, Text 17 | 18 | import pandas 19 | 20 | from picatrix.lib import framework, manager, state 21 | 22 | 23 | @framework.picatrix_magic 24 | def picatrixhelpers(data: Optional[Text] = '') -> pandas.DataFrame: 25 | """Provides information about registered Picatrix helpers. 26 | 27 | Args: 28 | data (str): If empty an overview of all registered helpers is provided, 29 | otherwise the help message of a particular helper is provided. 30 | 31 | Returns: 32 | A pandas DataFrame that contains the names and basic information about 33 | every registered helper or information about a single helper if 34 | the data string is provided. 35 | """ 36 | info_df = manager.MagicManager.get_helper_info(as_pandas=True) 37 | if not data: 38 | return info_df 39 | 40 | helper_obj = manager.MagicManager.get_helper(data) 41 | if not helper_obj: 42 | return pandas.DataFrame() 43 | 44 | return info_df[info_df.name == data.strip()] 45 | 46 | 47 | @framework.picatrix_magic 48 | def picatrixmagics(data: Optional[Text] = '') -> pandas.DataFrame: 49 | """Provides information about registered Picatrix magics. 50 | 51 | Args: 52 | data (str): If empty an overview of all registered magics is provided, 53 | otherwise the help message of a particular magic is provided. 54 | 55 | Returns: 56 | A pandas DataFrame that contains the names and basic information about 57 | every registered magic or information about a single magic if 58 | the data string is provided. 59 | """ 60 | if not data: 61 | return manager.MagicManager.get_magic_info(as_pandas=True) 62 | 63 | magic_obj = manager.MagicManager.get_magic(data) 64 | 65 | if not magic_obj: 66 | return pandas.DataFrame() 67 | 68 | description = magic_obj.__doc__.split('\n')[0] 69 | return pandas.DataFrame( 70 | [ 71 | { 72 | 'name': magic_obj.magic_name, 73 | 'description': description, 74 | 'function': '{0:s}_func'.format(magic_obj.magic_name), 75 | 'help': magic_obj.argument_parser.format_help() 76 | } 77 | ]) 78 | 79 | 80 | # pylint: disable=unused-argument 81 | @framework.picatrix_magic 82 | def last_output(data: Optional[Text] = '') -> Any: 83 | """Returns the last output from a magic that was executed. 84 | 85 | Args: 86 | data (str): optional string that does nothing. 87 | 88 | Returns: 89 | The last output from a magic that was run. 90 | """ 91 | state_obj = state.state() 92 | output = state_obj.last_output 93 | 94 | return output 95 | -------------------------------------------------------------------------------- /picatrix/magics/common_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Copyright 2020 Google LLC 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # https://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """Tests for the picatrix framework.""" 16 | from typing import Optional, Text 17 | 18 | from picatrix.lib import framework, manager 19 | from picatrix.magics import common 20 | 21 | 22 | def my_very_own_test_magic(data: Text, stuff: Optional[int] = 20) -> Text: 23 | """This is a magic that is used for testing. 24 | 25 | Args: 26 | data (str): This is a string. 27 | stuff (int): And this is a number. 28 | 29 | Returns: 30 | str: A string that combines the two parameters together. 31 | """ 32 | return f'{data.strip()} - {stuff}' 33 | 34 | 35 | def test_registration(): 36 | """Test the magic decorator.""" 37 | manager.MagicManager.clear_magics() 38 | _ = framework.picatrix_magic(my_very_own_test_magic) 39 | magic = framework.picatrix_magic(common.picatrixmagics) 40 | 41 | df = magic(line='') 42 | assert len(df) == 2 43 | assert set(df.columns) == set(manager.MagicManager.MAGICS_DF_COLUMNS) 44 | assert set(df.name.unique()) == set( 45 | ['my_very_own_test_magic', 'picatrixmagics']) 46 | 47 | magic = framework.picatrix_magic(common.last_output) 48 | same_df = magic(line='') 49 | assert same_df.equals(df) 50 | -------------------------------------------------------------------------------- /picatrix/notebook_init.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of 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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Initialization module for picatrix magics. 15 | 16 | When starting a new notebook using picatrix it is enough to do 17 | 18 | from picatrix import notebook_init 19 | 20 | notebook_init.init() 21 | 22 | And that would register magics and initialize the notebook to be able 23 | to take advantage of picatrix magics and helper functions. 24 | """ 25 | # pylint: disable=unused-import 26 | from picatrix import helpers, magics 27 | from picatrix.lib import state 28 | 29 | 30 | def init(): 31 | """Initialize the notebook.""" 32 | # Initialize the state object. 33 | _ = state.state() 34 | -------------------------------------------------------------------------------- /picatrix/version.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google Inc. All rights reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://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, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Version information for picatrix.""" 15 | 16 | __version__ = '20210519' 17 | 18 | 19 | def get_version(): 20 | """Returns the version information.""" 21 | return __version__ 22 | -------------------------------------------------------------------------------- /prepare-picatrix-runtime.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Initialize Picatrix runtime with right set of Jupyter extensions. 3 | 4 | # Exit on error. Append "|| true" if you expect an error. 5 | set -o errexit 6 | # Exit on error inside any functions or subshells. 7 | set -o errtrace 8 | # Do not allow use of undefined vars. Use ${VAR:-} to use an undefined VAR 9 | set -o nounset 10 | # Catch the error in case mysqldump fails (but gzip succeeds) in `mysqldump |gzip` 11 | set -o pipefail 12 | 13 | jupyter serverextension enable --py jupyter_http_over_ws 14 | jupyter nbextension enable --py widgetsnbextension --sys-prefix 15 | jupyter contrib nbextension install --user 16 | jupyter nbextensions_configurator enable --user 17 | jupyter nbextension enable --py --user ipyaggrid 18 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | dfdatetime>=20210509 2 | google-auth-oauthlib>=0.4.1 3 | google-auth>=1.17.2 4 | ipyaggrid>=0.2.1 5 | ipython>=5.5.0 6 | numpy>=1.19.5 7 | pandas>=1.1.3 8 | timesketch-api-client>=20201130 9 | timesketch-import-client>=20200910 10 | typing-extensions>=3.7.4.3 11 | docstring-parser>=0.12 12 | -------------------------------------------------------------------------------- /requirements_dev.txt: -------------------------------------------------------------------------------- 1 | coverage~=5.0.2 2 | isort~=5.8.0 3 | mock~=2.0.0 4 | pre-commit~=2.15.0 5 | pylint~=2.6.0 6 | pytest~=6.2.5 7 | yapf~=0.31.0 8 | -------------------------------------------------------------------------------- /requirements_runtime.txt: -------------------------------------------------------------------------------- 1 | altair>=4.1.0 2 | ipyaggrid>=0.2.1 3 | ipython>=5.5.0 4 | ipywidgets>=5.1.1 5 | jupyter-contrib-nbextensions>=0.5.1 6 | jupyter-http-over-ws>=0.0.8 7 | jupyter>=1.0.0 8 | matplotlib>=2.2.0 9 | numpy>=1.19.0 10 | pandas>=1.1.3 11 | scikit-learn>=1.0 12 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # Copyright 2020 Google Inc. All rights reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | """This is the setup file for the project.""" 16 | from typing import List, Text 17 | 18 | from setuptools import find_packages, setup 19 | 20 | 21 | def from_file(name: Text) -> List[Text]: 22 | """Read dependencies from requirements.txt file.""" 23 | with open(name, "r", encoding="utf-8") as f: 24 | return f.read().splitlines() 25 | 26 | 27 | long_description = ( 28 | "picatrix - a framework to assist security analysts using " 29 | "Colab or Jupyter to perform forensic investigations.") 30 | 31 | setup( 32 | name="picatrix", 33 | version="20210519", 34 | description="Picatrix IPython Helpers", 35 | long_description=long_description, 36 | license="Apache License, Version 2.0", 37 | url="https://github.com/google/picatrix/", 38 | maintainer="Picatrix development team", 39 | maintainer_email="picatrix-developers@googlegroups.com", 40 | classifiers=[ 41 | "Development Status :: 4 - Beta", 42 | "Environment :: Console", 43 | "Operating System :: OS Independent", 44 | "Programming Language :: Python", 45 | ], 46 | packages=find_packages(), 47 | include_package_data=True, 48 | zip_safe=False, 49 | install_requires=from_file("requirements.txt"), 50 | tests_require=from_file("requirements_dev.txt"), 51 | extras_require={ 52 | "runtime": from_file("requirements_runtime.txt"), 53 | }) 54 | --------------------------------------------------------------------------------