pax_global_header00006660000000000000000000000064152557723550014532gustar00rootroot0000000000000052 comment=fd3bc819bd7153dc5e63aad89150c93d29e827bf cthoyt-pystow-54a586e/000077500000000000000000000000001525577235500147525ustar00rootroot00000000000000cthoyt-pystow-54a586e/.cruft.json000066400000000000000000000021121525577235500170420ustar00rootroot00000000000000{ "template": "https://github.com/cthoyt/cookiecutter-snekpack", "commit": "3257f18b9b5dc6922830dea64f9e0ab8b42a40e4", "checkout": null, "context": { "cookiecutter": { "package_name": "pystow", "package_name_stylized": "PyStow", "short_description": "Easily pick a place to store data for your Python code", "author_name": "Charles Tapley Hoyt", "author_github": "cthoyt", "author_email": "cthoyt@gmail.com", "github_organization_name": "cthoyt", "github_repository_name": "pystow", "command_line_interface": true, "gitlab": false, "runner": "tox", "__runner": "tox -e", "__runner_uv": "--with tox-uv tox -e", "__runner_pip": "tox tox-uv", "__runner_install_uv": "uv tool install tox --with tox-uv", "__runner_install_pip": "python3 -m pip install tox tox-uv", "__runner_tests": "py", "__gh_slug": "cthoyt/pystow", "_template": "https://github.com/cthoyt/cookiecutter-snekpack", "_commit": "3257f18b9b5dc6922830dea64f9e0ab8b42a40e4" } }, "directory": null } cthoyt-pystow-54a586e/.gitattributes000066400000000000000000000001541525577235500176450ustar00rootroot00000000000000# Preserve LF line endings for test resources tests/resources/*.txt text eol=lf tests/resources/*.gz binary cthoyt-pystow-54a586e/.github/000077500000000000000000000000001525577235500163125ustar00rootroot00000000000000cthoyt-pystow-54a586e/.github/CODE_OF_CONDUCT.md000066400000000000000000000125511525577235500211150ustar00rootroot00000000000000# Contributor Covenant Code of Conduct ## Our Pledge We as members, contributors, and leaders pledge to make participation in our community a harassment-free experience for everyone, regardless of age, body size, visible or invisible disability, ethnicity, sex characteristics, gender identity and expression, level of experience, education, socio-economic status, nationality, personal appearance, race, caste, color, religion, or sexual identity and orientation. We pledge to act and interact in ways that contribute to an open, welcoming, diverse, inclusive, and healthy community. ## Our Standards Examples of behavior that contributes to a positive environment for our community include: - Demonstrating empathy and kindness toward other people - Being respectful of differing opinions, viewpoints, and experiences - Giving and gracefully accepting constructive feedback - Accepting responsibility and apologizing to those affected by our mistakes, and learning from the experience - Focusing on what is best not just for us as individuals, but for the overall community Examples of unacceptable behavior include: - The use of sexualized language or imagery, and sexual attention or advances of any kind - Trolling, insulting or derogatory comments, and personal or political attacks - Public or private harassment - Publishing others' private information, such as a physical or email address, without their explicit permission - Other conduct which could reasonably be considered inappropriate in a professional setting ## Enforcement Responsibilities Community leaders are responsible for clarifying and enforcing our standards of acceptable behavior and will take appropriate and fair corrective action in response to any behavior that they deem inappropriate, threatening, offensive, or harmful. Community leaders have the right and responsibility to remove, edit, or reject comments, commits, code, wiki edits, issues, and other contributions that are not aligned to this Code of Conduct, and will communicate reasons for moderation decisions when appropriate. ## Scope This Code of Conduct applies within all community spaces, and also applies when an individual is officially representing the community in public spaces. Examples of representing our community include using an official e-mail address, posting via an official social media account, or acting as an appointed representative at an online or offline event. ## Enforcement Instances of abusive, harassing, or otherwise unacceptable behavior may be reported to the community leaders responsible for enforcement at cthoyt@gmail.com. All complaints will be reviewed and investigated promptly and fairly. All community leaders are obligated to respect the privacy and security of the reporter of any incident. ## Enforcement Guidelines Community leaders will follow these Community Impact Guidelines in determining the consequences for any action they deem in violation of this Code of Conduct: ### 1. Correction **Community Impact**: Use of inappropriate language or other behavior deemed unprofessional or unwelcome in the community. **Consequence**: A private, written warning from community leaders, providing clarity around the nature of the violation and an explanation of why the behavior was inappropriate. A public apology may be requested. ### 2. Warning **Community Impact**: A violation through a single incident or series of actions. **Consequence**: A warning with consequences for continued behavior. No interaction with the people involved, including unsolicited interaction with those enforcing the Code of Conduct, for a specified period of time. This includes avoiding interactions in community spaces as well as external channels like social media. Violating these terms may lead to a temporary or permanent ban. ### 3. Temporary Ban **Community Impact**: A serious violation of community standards, including sustained inappropriate behavior. **Consequence**: A temporary ban from any sort of interaction or public communication with the community for a specified period of time. No public or private interaction with the people involved, including unsolicited interaction with those enforcing the Code of Conduct, is allowed during this period. Violating these terms may lead to a permanent ban. ### 4. Permanent Ban **Community Impact**: Demonstrating a pattern of violation of community standards, including sustained inappropriate behavior, harassment of an individual, or aggression toward or disparagement of classes of individuals. **Consequence**: A permanent ban from any sort of public interaction within the community. ## Attribution This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 2.1, available at [https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1]. Community Impact Guidelines were inspired by [Mozilla's code of conduct enforcement ladder][Mozilla CoC]. For answers to common questions about this code of conduct, see the FAQ at [https://www.contributor-covenant.org/faq][FAQ]. Translations are available at [https://www.contributor-covenant.org/translations][translations]. [homepage]: https://www.contributor-covenant.org [v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html [Mozilla CoC]: https://github.com/mozilla/diversity [FAQ]: https://www.contributor-covenant.org/faq [translations]: https://www.contributor-covenant.org/translations cthoyt-pystow-54a586e/.github/CONTRIBUTING.md000066400000000000000000000120301525577235500205370ustar00rootroot00000000000000# Contributing Contributions to this repository are welcomed and encouraged. ## Code Contribution This project uses the [GitHub Flow](https://guides.github.com/introduction/flow) model for code contributions. Follow these steps: 1. [Create a fork](https://help.github.com/articles/fork-a-repo) of the upstream repository at [`cthoyt/pystow`](https://github.com/cthoyt/pystow) on your GitHub account (or in one of your organizations) 2. [Clone your fork](https://docs.github.com/en/repositories/creating-and-managing-repositories/cloning-a-repository) with `git clone https://github.com//pystow.git` 3. Make and commit changes to your fork with `git commit` 4. Push changes to your fork with `git push` 5. Repeat steps 3 and 4 as needed 6. Submit a pull request back to the upstream repository ### Merge Model This repository uses [squash merges](https://docs.github.com/en/github/collaborating-with-pull-requests/incorporating-changes-from-a-pull-request/about-pull-request-merges#squash-and-merge-your-pull-request-commits) to group all related commits in a given pull request into a single commit upon acceptance and merge into the main branch. This has several benefits: 1. Keeps the commit history on the main branch focused on high-level narrative 2. Enables people to make lots of small commits without worrying about muddying up the commit history 3. Commits correspond 1-to-1 with pull requests ### Code Style This project uses `tox` for running code quality checks. Start by installing it with `pip install tox tox-uv`. This project encourages the use of optional static typing. It uses [`mypy`](http://mypy-lang.org/) as a type checker. You can check if your code passes `mypy` with `tox -e mypy`. This project uses [`ruff`](https://docs.astral.sh/ruff/) to automatically enforce a consistent code style. You can apply `ruff format` and other pre-configured formatters with `tox -e format`. This project uses [`ruff`](https://docs.astral.sh/ruff/) and several plugins for additional checks of documentation style, security issues, good variable nomenclature, and more (see `pyproject.toml` for a list of Ruff plugins). You can check if your code passes `ruff check` with `tox -e lint`. Each of these checks are run on each commit using GitHub Actions as a continuous integration service. Passing all of them is required for accepting a contribution. If you're unsure how to address the feedback from one of these tools, please say so either in the description of your pull request or in a comment, and we will help you. ### Logging Python's builtin `print()` should not be used (except when writing to files), it's checked by the [`flake8-print` (T20)](https://docs.astral.sh/ruff/rules/#flake8-print-t20) plugin to `ruff`. If you're in a command line setting or `main()` function for a module, you can use `click.echo()`. Otherwise, you can use the builtin `logging` module by adding `logger = logging.getLogger(__name__)` below the imports at the top of your file. ### Documentation All public functions (i.e., not starting with an underscore `_`) must be documented using the [sphinx documentation format](https://sphinx-rtd-tutorial.readthedocs.io/en/latest/docstrings.html#the-sphinx-docstring-format). The [`darglint2`](https://github.com/akaihola/darglint2) tool reports on functions that are not fully documented. This project uses [`sphinx`](https://www.sphinx-doc.org) to automatically build documentation into a narrative structure. You can check that the documentation builds properly in an isolated environment with `tox -e docs-test` and actually build it locally with `tox -e docs`. ### Testing Functions in this repository should be unit tested. These can either be written using the `unittest` framework in the `tests/` directory or as embedded doctests. You can check that the unit tests pass with `tox -e py` and that the doctests pass with `tox -e doctests`. These tests are required to pass for accepting a contribution. ### Syncing your fork If other code is updated before your contribution gets merged, you might need to resolve conflicts against the main branch. After cloning, you should add the upstream repository with ```shell $ git remote add cthoyt https://github.com/cthoyt/pystow.git ``` Then, you can merge upstream code into your branch. You can also use the GitHub UI to do this by following [this tutorial](https://docs.github.com/en/github/collaborating-with-pull-requests/working-with-forks/syncing-a-fork). ### Python Version Compatibility This project aims to support all versions of Python that have not passed their end-of-life dates. After end-of-life, the version will be removed from the Trove qualifiers in the `pyproject.toml` and from the GitHub Actions testing configuration. See https://endoflife.date/python for a timeline of Python release and end-of-life dates. ## Acknowledgements These code contribution guidelines are derived from the [cthoyt/cookiecutter-snekpack](https://github.com/cthoyt/cookiecutter-snekpack) Python package template. They're free to reuse and modify as long as they're properly acknowledged. cthoyt-pystow-54a586e/.github/FUNDING.yml000066400000000000000000000002701525577235500201260ustar00rootroot00000000000000# https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/displaying-a-sponsor-button-in-your-repository github: - cthoyt cthoyt-pystow-54a586e/.github/PULL_REQUEST_TEMPLATE.md000066400000000000000000000011241525577235500221110ustar00rootroot00000000000000 ## Summary cthoyt-pystow-54a586e/.github/codecov.yml000066400000000000000000000001631525577235500204570ustar00rootroot00000000000000# see https://docs.codecov.com/v4.6/docs/codecov-yaml ignore: - "src/pystow/__main__.py" - "src/pystow/cli.py" cthoyt-pystow-54a586e/.github/workflows/000077500000000000000000000000001525577235500203475ustar00rootroot00000000000000cthoyt-pystow-54a586e/.github/workflows/tests.yml000066400000000000000000000074201525577235500222370ustar00rootroot00000000000000# This file configures the continuous integration (CI) system on GitHub. # Introductory materials can be found here: https://docs.github.com/en/actions/learn-github-actions/understanding-github-actions. # Documentation for editing this file can be found here: https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions name: Tests # by default, give the GITHUB_TOKEN no permissions # See https://docs.github.com/en/actions/writing-workflows/choosing-what-your-workflow-does/controlling-permissions-for-github_token permissions: { } on: push: branches: [ main ] pull_request: branches: [ main ] jobs: lint: name: Code Quality permissions: # give only read-only access to the contents of the repository # this is the only permission this job requires, so keep it to the least privilege # i.e., not to issues, discussions, actions, etc. contents: read runs-on: ubuntu-latest strategy: matrix: python-version: [ "3.14", "3.11" ] tox-command: [ "lint", "pyroma", "mypy" ] steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: persist-credentials: false - name: "Install uv" uses: astral-sh/setup-uv@cec208311dfd045dd5311c1add060b2062131d57 # v8.0.0 with: enable-cache: true cache-dependency-glob: "pyproject.toml" - name: "Run command" run: | uvx -p ${{ matrix.python-version }} --with tox-uv tox -e ${{ matrix.tox-command }} docs: name: Documentation permissions: contents: read runs-on: ubuntu-latest strategy: matrix: python-version: [ "3.14" ] steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: persist-credentials: false - name: "Install uv" uses: astral-sh/setup-uv@cec208311dfd045dd5311c1add060b2062131d57 # v8.0.0 with: enable-cache: true cache-dependency-glob: "pyproject.toml" - name: Install dependencies run: | sudo apt-get install graphviz - name: Check RST conformity with doc8 run: uvx -p ${{ matrix.python-version }} --with tox-uv tox -e doc8 - name: Check docstring coverage run: uvx -p ${{ matrix.python-version }} --with tox-uv tox -e docstr-coverage - name: Check documentation build with Sphinx run: uvx -p ${{ matrix.python-version }} --with tox-uv tox -e docs-test # Check all Markdown files are properly formatted # inspired by https://github.com/astral-sh/uv/blob/98523e2014e9a5c69706623344026d76296e178f/.github/workflows/ci.yml#L67C1-L70C61 - name: Check markdown formatting run: | npx --yes prettier --prose-wrap always --check "**/*.md" tests: name: Tests permissions: contents: read runs-on: ${{ matrix.os }} strategy: matrix: os: [ ubuntu-latest, windows-latest ] python-version: [ "3.14", "3.11" ] steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: persist-credentials: false - name: "Install uv" uses: astral-sh/setup-uv@cec208311dfd045dd5311c1add060b2062131d57 # v8.0.0 with: enable-cache: true cache-dependency-glob: "pyproject.toml" - name: Test with pytest and generate coverage file run: uvx -p ${{ matrix.python-version }} --with tox-uv tox -e py - name: Run doctests run: uvx -p ${{ matrix.python-version }} --with tox-uv tox -e doctests - name: Upload coverage report to codecov uses: astral-sh/setup-uv@cec208311dfd045dd5311c1add060b2062131d57 # v8.0.0 if: success() with: file: coverage.xml cthoyt-pystow-54a586e/.gitignore000066400000000000000000000402141525577235500167430ustar00rootroot00000000000000# Created by https://www.toptal.com/developers/gitignore/api/macos,linux,windows,python,jupyternotebooks,jetbrains,pycharm,vim,emacs,visualstudiocode,visualstudio # Edit at https://www.toptal.com/developers/gitignore?templates=macos,linux,windows,python,jupyternotebooks,jetbrains,pycharm,vim,emacs,visualstudiocode,visualstudio ### Emacs ### # -*- mode: gitignore; -*- *~ \#*\# /.emacs.desktop /.emacs.desktop.lock *.elc auto-save-list tramp .\#* # Org-mode .org-id-locations *_archive # flymake-mode *_flymake.* # eshell files /eshell/history /eshell/lastdir # elpa packages /elpa/ # reftex files *.rel # AUCTeX auto folder /auto/ # cask packages .cask/ dist/ # Flycheck flycheck_*.el # server auth directory /server/ # projectiles files .projectile # directory configuration .dir-locals.el # network security /network-security.data ### JetBrains ### # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 # User-specific stuff .idea/**/workspace.xml .idea/**/tasks.xml .idea/**/usage.statistics.xml .idea/**/dictionaries .idea/**/shelf # AWS User-specific .idea/**/aws.xml # Generated files .idea/**/contentModel.xml # Sensitive or high-churn files .idea/**/dataSources/ .idea/**/dataSources.ids .idea/**/dataSources.local.xml .idea/**/sqlDataSources.xml .idea/**/dynamic.xml .idea/**/uiDesigner.xml .idea/**/dbnavigator.xml # Gradle .idea/**/gradle.xml .idea/**/libraries # Gradle and Maven with auto-import # When using Gradle or Maven with auto-import, you should exclude module files, # since they will be recreated, and may cause churn. Uncomment if using # auto-import. # .idea/artifacts # .idea/compiler.xml # .idea/jarRepositories.xml # .idea/modules.xml # .idea/*.iml # .idea/modules # *.iml # *.ipr # CMake cmake-build-*/ # Mongo Explorer plugin .idea/**/mongoSettings.xml # File-based project format *.iws # IntelliJ out/ # mpeltonen/sbt-idea plugin .idea_modules/ # JIRA plugin atlassian-ide-plugin.xml # Cursive Clojure plugin .idea/replstate.xml # SonarLint plugin .idea/sonarlint/ # Crashlytics plugin (for Android Studio and IntelliJ) com_crashlytics_export_strings.xml crashlytics.properties crashlytics-build.properties fabric.properties # Editor-based Rest Client .idea/httpRequests # Android studio 3.1+ serialized cache file .idea/caches/build_file_checksums.ser ### JetBrains Patch ### # Comment Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-215987721 # *.iml # modules.xml # .idea/misc.xml # *.ipr # Sonarlint plugin # https://plugins.jetbrains.com/plugin/7973-sonarlint .idea/**/sonarlint/ # SonarQube Plugin # https://plugins.jetbrains.com/plugin/7238-sonarqube-community-plugin .idea/**/sonarIssues.xml # Markdown Navigator plugin # https://plugins.jetbrains.com/plugin/7896-markdown-navigator-enhanced .idea/**/markdown-navigator.xml .idea/**/markdown-navigator-enh.xml .idea/**/markdown-navigator/ # Cache file creation bug # See https://youtrack.jetbrains.com/issue/JBR-2257 .idea/$CACHE_FILE$ # CodeStream plugin # https://plugins.jetbrains.com/plugin/12206-codestream .idea/codestream.xml ### JupyterNotebooks ### # gitignore template for Jupyter Notebooks # website: http://jupyter.org/ .ipynb_checkpoints */.ipynb_checkpoints/* # IPython profile_default/ ipython_config.py # Remove previous ipynb_checkpoints # git rm -r .ipynb_checkpoints/ ### Linux ### # temporary files which can be created if a process still has a handle open of a deleted file .fuse_hidden* # KDE directory preferences .directory # Linux trash folder which might appear on any partition or disk .Trash-* # .nfs files are created when an open file is removed but is still being accessed .nfs* ### macOS ### # General .DS_Store .AppleDouble .LSOverride # Icon must end with two \r Icon # Thumbnails ._* # Files that might appear in the root of a volume .DocumentRevisions-V100 .fseventsd .Spotlight-V100 .TemporaryItems .Trashes .VolumeIcon.icns .com.apple.timemachine.donotpresent # Directories potentially created on remote AFP share .AppleDB .AppleDesktop Network Trash Folder Temporary Items .apdisk ### PyCharm ### # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 # User-specific stuff # AWS User-specific # Generated files # Sensitive or high-churn files # Gradle # Gradle and Maven with auto-import # When using Gradle or Maven with auto-import, you should exclude module files, # since they will be recreated, and may cause churn. Uncomment if using # auto-import. # .idea/artifacts # .idea/compiler.xml # .idea/jarRepositories.xml # .idea/modules.xml # .idea/*.iml # .idea/modules # *.iml # *.ipr # CMake # Mongo Explorer plugin # File-based project format # IntelliJ # mpeltonen/sbt-idea plugin # JIRA plugin # Cursive Clojure plugin # SonarLint plugin # Crashlytics plugin (for Android Studio and IntelliJ) # Editor-based Rest Client # Android studio 3.1+ serialized cache file ### PyCharm Patch ### # Comment Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-215987721 # *.iml # modules.xml # .idea/misc.xml # *.ipr # Sonarlint plugin # https://plugins.jetbrains.com/plugin/7973-sonarlint # SonarQube Plugin # https://plugins.jetbrains.com/plugin/7238-sonarqube-community-plugin # Markdown Navigator plugin # https://plugins.jetbrains.com/plugin/7896-markdown-navigator-enhanced # Cache file creation bug # See https://youtrack.jetbrains.com/issue/JBR-2257 # CodeStream plugin # https://plugins.jetbrains.com/plugin/12206-codestream ### Python ### # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so # Distribution / packaging .Python build/ develop-eggs/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .nox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *.cover *.py,cover .hypothesis/ .pytest_cache/ cover/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py db.sqlite3 db.sqlite3-journal # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ docs/build docs/source/api # PyBuilder .pybuilder/ target/ # Jupyter Notebook # IPython # pyenv # For a library or package, you might want to ignore these files since the code is # intended to run in multiple environments; otherwise, check them in: # .python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. # However, in case of collaboration, if having platform-specific dependencies or dependencies # having no cross-platform support, pipenv may install dependencies that don't work, or not # install all needed dependencies. #Pipfile.lock # poetry # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. # This is especially recommended for binary packages to ensure reproducibility, and is more # commonly ignored for libraries. # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control #poetry.lock # PEP 582; used by e.g. github.com/David-OConnor/pyflow __pypackages__/ # Celery stuff celerybeat-schedule celerybeat.pid # SageMath parsed files *.sage.py # Environments .env .venv env/ venv/ ENV/ env.bak/ venv.bak/ # Spyder project settings .spyderproject .spyproject # Rope project settings .ropeproject # mkdocs documentation /site # mypy .mypy_cache/ .dmypy.json dmypy.json # Pyre type checker .pyre/ # pytype static type analyzer .pytype/ # Cython debug symbols cython_debug/ # PyCharm # JetBrains specific template is maintainted in a separate JetBrains.gitignore that can # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ ### Vim ### # Swap [._]*.s[a-v][a-z] !*.svg # comment out if you don't need vector files [._]*.sw[a-p] [._]s[a-rt-v][a-z] [._]ss[a-gi-z] [._]sw[a-p] # Session Session.vim Sessionx.vim # Temporary .netrwhist # Auto-generated tag files tags # Persistent undo [._]*.un~ ### VisualStudioCode ### .vscode/* !.vscode/settings.json !.vscode/tasks.json !.vscode/launch.json !.vscode/extensions.json !.vscode/*.code-snippets # Local History for Visual Studio Code .history/ # Built Visual Studio Code Extensions *.vsix ### VisualStudioCode Patch ### # Ignore all local history of files .history .ionide # Support for Project snippet scope ### Windows ### # Windows thumbnail cache files Thumbs.db Thumbs.db:encryptable ehthumbs.db ehthumbs_vista.db # Dump file *.stackdump # Folder config file [Dd]esktop.ini # Recycle Bin used on file shares $RECYCLE.BIN/ # Windows Installer files *.cab *.msi *.msix *.msm *.msp # Windows shortcuts *.lnk ### VisualStudio ### ## Ignore Visual Studio temporary files, build results, and ## files generated by popular Visual Studio add-ons. ## ## Get latest from https://github.com/github/gitignore/blob/main/VisualStudio.gitignore # User-specific files *.rsuser *.suo *.user *.userosscache *.sln.docstates # User-specific files (MonoDevelop/Xamarin Studio) *.userprefs # Mono auto generated files mono_crash.* # Build results [Dd]ebug/ [Dd]ebugPublic/ [Rr]elease/ [Rr]eleases/ x64/ x86/ [Ww][Ii][Nn]32/ [Aa][Rr][Mm]/ [Aa][Rr][Mm]64/ bld/ [Bb]in/ [Oo]bj/ [Ll]og/ [Ll]ogs/ # Visual Studio 2015/2017 cache/options directory .vs/ # Uncomment if you have tasks that create the project's static files in wwwroot #wwwroot/ # Visual Studio 2017 auto generated files Generated\ Files/ # MSTest test Results [Tt]est[Rr]esult*/ [Bb]uild[Ll]og.* # NUnit *.VisualState.xml TestResult.xml nunit-*.xml # Build Results of an ATL Project [Dd]ebugPS/ [Rr]eleasePS/ dlldata.c # Benchmark Results BenchmarkDotNet.Artifacts/ # .NET Core project.lock.json project.fragment.lock.json artifacts/ # ASP.NET Scaffolding ScaffoldingReadMe.txt # StyleCop StyleCopReport.xml # Files built by Visual Studio *_i.c *_p.c *_h.h *.ilk *.meta *.obj *.iobj *.pch *.pdb *.ipdb *.pgc *.pgd *.rsp *.sbr *.tlb *.tli *.tlh *.tmp *.tmp_proj *_wpftmp.csproj *.tlog *.vspscc *.vssscc .builds *.pidb *.svclog *.scc # Chutzpah Test files _Chutzpah* # Visual C++ cache files ipch/ *.aps *.ncb *.opendb *.opensdf *.sdf *.cachefile *.VC.db *.VC.VC.opendb # Visual Studio profiler *.psess *.vsp *.vspx *.sap # Visual Studio Trace Files *.e2e # TFS 2012 Local Workspace $tf/ # Guidance Automation Toolkit *.gpState # ReSharper is a .NET coding add-in _ReSharper*/ *.[Rr]e[Ss]harper *.DotSettings.user # TeamCity is a build add-in _TeamCity* # DotCover is a Code Coverage Tool *.dotCover # AxoCover is a Code Coverage Tool .axoCover/* !.axoCover/settings.json # Coverlet is a free, cross platform Code Coverage Tool coverage*.json coverage*.xml coverage*.info # Visual Studio code coverage results *.coverage *.coveragexml # NCrunch _NCrunch_* .*crunch*.local.xml nCrunchTemp_* # MightyMoose *.mm.* AutoTest.Net/ # Web workbench (sass) .sass-cache/ # Installshield output folder [Ee]xpress/ # DocProject is a documentation generator add-in DocProject/buildhelp/ DocProject/Help/*.HxT DocProject/Help/*.HxC DocProject/Help/*.hhc DocProject/Help/*.hhk DocProject/Help/*.hhp DocProject/Help/Html2 DocProject/Help/html # Click-Once directory publish/ # Publish Web Output *.[Pp]ublish.xml *.azurePubxml # Note: Comment the next line if you want to checkin your web deploy settings, # but database connection strings (with potential passwords) will be unencrypted *.pubxml *.publishproj # Microsoft Azure Web App publish settings. Comment the next line if you want to # checkin your Azure Web App publish settings, but sensitive information contained # in these scripts will be unencrypted PublishScripts/ # NuGet Packages *.nupkg # NuGet Symbol Packages *.snupkg # The packages folder can be ignored because of Package Restore **/[Pp]ackages/* # except build/, which is used as an MSBuild target. !**/[Pp]ackages/build/ # Uncomment if necessary however generally it will be regenerated when needed #!**/[Pp]ackages/repositories.config # NuGet v3's project.json files produces more ignorable files *.nuget.props *.nuget.targets # Microsoft Azure Build Output csx/ *.build.csdef # Microsoft Azure Emulator ecf/ rcf/ # Windows Store app package directories and files AppPackages/ BundleArtifacts/ Package.StoreAssociation.xml _pkginfo.txt *.appx *.appxbundle *.appxupload # Visual Studio cache files # files ending in .cache can be ignored *.[Cc]ache # but keep track of directories ending in .cache !?*.[Cc]ache/ # Others ClientBin/ ~$* *.dbmdl *.dbproj.schemaview *.jfm *.pfx *.publishsettings orleans.codegen.cs # Including strong name files can present a security risk # (https://github.com/github/gitignore/pull/2483#issue-259490424) #*.snk # Since there are multiple workflows, uncomment next line to ignore bower_components # (https://github.com/github/gitignore/pull/1529#issuecomment-104372622) #bower_components/ # RIA/Silverlight projects Generated_Code/ # Backup & report files from converting an old project file # to a newer Visual Studio version. Backup files are not needed, # because we have git ;-) _UpgradeReport_Files/ Backup*/ UpgradeLog*.XML UpgradeLog*.htm ServiceFabricBackup/ *.rptproj.bak # SQL Server files *.mdf *.ldf *.ndf # Business Intelligence projects *.rdl.data *.bim.layout *.bim_*.settings *.rptproj.rsuser *- [Bb]ackup.rdl *- [Bb]ackup ([0-9]).rdl *- [Bb]ackup ([0-9][0-9]).rdl # Microsoft Fakes FakesAssemblies/ # GhostDoc plugin setting file *.GhostDoc.xml # Node.js Tools for Visual Studio .ntvs_analysis.dat node_modules/ # Visual Studio 6 build log *.plg # Visual Studio 6 workspace options file *.opt # Visual Studio 6 auto-generated workspace file (contains which files were open etc.) *.vbw # Visual Studio 6 auto-generated project file (contains which files were open etc.) *.vbp # Visual Studio 6 workspace and project file (working project files containing files to include in project) *.dsw *.dsp # Visual Studio 6 technical files # Visual Studio LightSwitch build output **/*.HTMLClient/GeneratedArtifacts **/*.DesktopClient/GeneratedArtifacts **/*.DesktopClient/ModelManifest.xml **/*.Server/GeneratedArtifacts **/*.Server/ModelManifest.xml _Pvt_Extensions # Paket dependency manager .paket/paket.exe paket-files/ # FAKE - F# Make .fake/ # CodeRush personal settings .cr/personal # Python Tools for Visual Studio (PTVS) *.pyc # Cake - Uncomment if you are using it # tools/** # !tools/packages.config # Tabs Studio *.tss # Telerik's JustMock configuration file *.jmconfig # BizTalk build output *.btp.cs *.btm.cs *.odx.cs *.xsd.cs # OpenCover UI analysis results OpenCover/ # Azure Stream Analytics local run output ASALocalRun/ # MSBuild Binary and Structured Log *.binlog # NVidia Nsight GPU debugger configuration file *.nvuser # MFractors (Xamarin productivity tool) working folder .mfractor/ # Local History for Visual Studio .localhistory/ # Visual Studio History (VSHistory) files .vshistory/ # BeatPulse healthcheck temp database healthchecksdb # Backup folder for Package Reference Convert tool in Visual Studio 2017 MigrationBackup/ # Ionide (cross platform F# VS Code tools) working folder .ionide/ # Fody - auto-generated XML schema FodyWeavers.xsd # VS Code files for those working on multiple tools *.code-workspace # Local History for Visual Studio Code # Windows Installer files from build outputs # JetBrains Rider *.sln.iml ### VisualStudio Patch ### # Additional files built by Visual Studio # End of https://www.toptal.com/developers/gitignore/api/macos,linux,windows,python,jupyternotebooks,jetbrains,pycharm,vim,emacs,visualstudiocode,visualstudio scratch/ tests/resources/test_1.db tests/resources/test_1.json.bz2 tests/resources/test_1.yaml tests/resources/test_1.pkl.gz cthoyt-pystow-54a586e/.readthedocs.yml000066400000000000000000000021051525577235500200360ustar00rootroot00000000000000# .readthedocs.yaml # Read the Docs configuration file # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details version: 2 sphinx: # Path to your Sphinx configuration file, required as of # https://about.readthedocs.com/blog/2024/12/deprecate-config-files-without-sphinx-or-mkdocs-config/ configuration: docs/source/conf.py # Set the version of Python and other tools you might need build: os: ubuntu-22.04 apt_packages: - graphviz tools: python: "3.12" # adapted from uv recipe at https://docs.readthedocs.io/en/stable/build-customization.html#install-dependencies-with-uv # and comment at https://github.com/readthedocs/readthedocs.org/issues/11289#issuecomment-2103832834 commands: - asdf plugin add uv - asdf install uv latest - asdf global uv latest - uv venv $READTHEDOCS_VIRTUALENV_PATH - VIRTUAL_ENV=$READTHEDOCS_VIRTUALENV_PATH uv pip install --group docs .[rdf,pandas,xml,aws,yaml,pydantic,requests,cli] - python -m sphinx -T -b html -d docs/_build/doctrees -D language=en docs/source $READTHEDOCS_OUTPUT/html cthoyt-pystow-54a586e/LICENSE000066400000000000000000000020641525577235500157610ustar00rootroot00000000000000MIT License Copyright (c) 2026 Charles Tapley Hoyt Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. cthoyt-pystow-54a586e/README.md000066400000000000000000000352271525577235500162420ustar00rootroot00000000000000

PyStow

Tests PyPI PyPI - Python Version PyPI - License Documentation Status Codecov status Cookiecutter template from @cthoyt Ruff Contributor Covenant DOI

👜 Easily pick a place to store data for your Python code ## 💪 Getting Started Get a directory for your application. ```python import pystow # Get a directory (as a pathlib.Path) for ~/.data/pykeen pykeen_directory = pystow.join("pykeen") # Get a subdirectory (as a pathlib.Path) for ~/.data/pykeen/experiments pykeen_experiments_directory = pystow.join("pykeen", "experiments") # You can go as deep as you want pykeen_deep_directory = pystow.join("pykeen", "experiments", "a", "b", "c") ``` If you reuse the same directory structure a lot, you can save them in a module: ```python import pystow pykeen_module = pystow.module("pykeen") # Access the module's directory with .base assert pystow.join("pykeen") == pystow.module("pykeen").base # Get a subdirectory (as a pathlib.Path) for ~/.data/pykeen/experiments pykeen_experiments_directory = pykeen_module.join("experiments") # You can go as deep as you want past the original "pykeen" module pykeen_deep_directory = pykeen_module.join("experiments", "a", "b", "c") ``` Get a file path for your application by adding the `name` keyword argument. This is made explicit so PyStow knows which parent directories to automatically create. This works with `pystow` or any module you create with `pystow.module`. ```python import pystow # Get a directory (as a pathlib.Path) for ~/.data/indra/database.tsv indra_database_path = pystow.join("indra", "database", name="database.tsv") ``` Ensure a file from the internet is available in your application's directory: ```python import pystow url = "https://raw.githubusercontent.com/pykeen/pykeen/master/src/pykeen/datasets/nations/test.txt" path = pystow.ensure("pykeen", "datasets", "nations", url=url) ``` Ensure a tabular data file from the internet and load it for usage (requires `pip install pandas`): ```python import pystow import pandas as pd url = "https://raw.githubusercontent.com/pykeen/pykeen/master/src/pykeen/datasets/nations/test.txt" df: pd.DataFrame = pystow.ensure_csv("pykeen", "datasets", "nations", url=url) ``` Ensure a comma-separated tabular data file from the internet and load it for usage (requires `pip install pandas`): ```python import pystow import pandas as pd url = "https://raw.githubusercontent.com/cthoyt/pystow/main/tests/resources/test_1.csv" df: pd.DataFrame = pystow.ensure_csv( "pykeen", "datasets", "nations", url=url, read_csv_kwargs=dict(sep=",") ) ``` Ensure a RDF file from the internet and load it for usage (requires `pip install rdflib`) ```python import pystow import rdflib url = "https://ftp.expasy.org/databases/rhea/rdf/rhea.rdf.gz" rdf_graph: rdflib.Graph = pystow.ensure_rdf("rhea", url=url) ``` Also see `pystow.ensure_excel()`, `pystow.ensure_rdf()`, `pystow.ensure_zip_df()`, and `pystow.ensure_tar_df()`. If your data comes with a lot of different files in an archive, you can ensure the archive is downloaded and get specific files from it: ```python import numpy as np import pystow url = "https://cloud.enterprise.informatik.uni-leipzig.de/index.php/s/LHPbMCre7SLqajB/download/MultiKE_D_Y_15K_V1.zip" # the path inside the archive to the file you want inner_path = "MultiKE/D_Y_15K_V1/721_5fold/1/20210219183115/ent_embeds.npy" with pystow.ensure_open_zip("kiez", url=url, inner_path=inner_path) as file: emb = np.load(file) ``` Also see `pystow.module.ensure_open_lzma()`, `pystow.module.ensure_open_tarfile()` and `pystow.module.ensure_open_gz()`. ## ⚙️️ Configuration By default, data is stored in the `$HOME/.data` directory. By default, the `` app will create the `$HOME/.data/` folder. If you want to use an alternate folder name to `.data` inside the home directory, you can set the `PYSTOW_NAME` environment variable. For example, if you set `PYSTOW_NAME=mydata`, then the following code for the `pykeen` app will create the `$HOME/mydata/pykeen/` directory: ```python import os import pystow # Only for demonstration purposes. You should set environment # variables either with your .bashrc or in the command line REPL. os.environ["PYSTOW_NAME"] = "mydata" # Get a directory (as a pathlib.Path) for ~/mydata/pykeen pykeen_directory = pystow.join("pykeen") ``` If you want to specify a completely custom directory that isn't relative to your home directory, you can set the `PYSTOW_HOME` environment variable. For example, if you set `PYSTOW_HOME=/usr/local/`, then the following code for the `pykeen` app will create the `/usr/local/pykeen/` directory: ```python import os import pystow # Only for demonstration purposes. You should set environment # variables either with your .bashrc or in the command line REPL. os.environ["PYSTOW_HOME"] = "/usr/local/" # Get a directory (as a pathlib.Path) for /usr/local/pykeen pykeen_directory = pystow.join("pykeen") ``` Note: if you set `PYSTOW_HOME`, then `PYSTOW_NAME` is disregarded. ### X Desktop Group (XDG) Compatibility While PyStow's main goal is to make application data less opaque and less hidden, some users might want to use the [XDG specifications](http://standards.freedesktop.org/basedir-spec/basedir-spec-latest.html) for storing their app data. If you set the environment variable `PYSTOW_USE_APPDIRS` to `true` or `True`, then the [`appdirs`](https://pypi.org/project/appdirs/) or [`platformdirs`](https://pypi.org/project/platformdirs/) package will be used to choose the base directory based on the `user data dir` option. This can still be overridden by `PYSTOW_HOME`. ## 🚀 Installation The most recent release can be installed from [PyPI](https://pypi.org/project/pystow/) with uv: ```console $ uv pip install pystow ``` or with pip: ```console $ python3 -m pip install pystow ``` The most recent code and data can be installed directly from GitHub with uv: ```console $ uv pip install git+https://github.com/cthoyt/pystow.git ``` or with pip: ```console $ python3 -m pip install git+https://github.com/cthoyt/pystow.git ``` ## 👐 Contributing Contributions, whether filing an issue, making a pull request, or forking, are appreciated. See [CONTRIBUTING.md](https://github.com/cthoyt/pystow/blob/master/.github/CONTRIBUTING.md) for more information on getting involved. ## 👋 Attribution ### ⚖️ License The code in this package is licensed under the MIT License. ### 🍪 Cookiecutter This package was created with [@audreyfeldroy](https://github.com/audreyfeldroy)'s [cookiecutter](https://github.com/cookiecutter/cookiecutter) package using [@cthoyt](https://github.com/cthoyt)'s [cookiecutter-snekpack](https://github.com/cthoyt/cookiecutter-snekpack) template. ## 🛠️ For Developers
See developer instructions The final section of the README is for if you want to get involved by making a code contribution. ### Development Installation To install in development mode, use the following: ```console $ git clone git+https://github.com/cthoyt/pystow.git $ cd pystow $ uv pip install -e . ``` Alternatively, install using pip: ```console $ python3 -m pip install -e . ``` ### Updating Package Boilerplate This project uses `cruft` to keep boilerplate (i.e., configuration, contribution guidelines, documentation configuration) up-to-date with the upstream cookiecutter package. Install cruft with either `uv tool install cruft` or `python3 -m pip install cruft` then run: ```console $ cruft update ``` More info on Cruft's update command is available [here](https://github.com/cruft/cruft?tab=readme-ov-file#updating-a-project). ### 🥼 Testing After cloning the repository and installing `tox` with `uv tool install tox --with tox-uv` or `python3 -m pip install tox tox-uv`, the unit tests in the `tests/` folder can be run reproducibly with: ```console $ tox -e py ``` Additionally, these tests are automatically re-run with each commit in a [GitHub Action](https://github.com/cthoyt/pystow/actions?query=workflow%3ATests). ### 📖 Building the Documentation The documentation can be built locally using the following: ```console $ git clone git+https://github.com/cthoyt/pystow.git $ cd pystow $ tox -e docs $ open docs/build/html/index.html ``` The documentation automatically installs the package as well as the `docs` extra specified in the [`pyproject.toml`](pyproject.toml). `sphinx` plugins like `texext` can be added there. Additionally, they need to be added to the `extensions` list in [`docs/source/conf.py`](docs/source/conf.py). The documentation can be deployed to [ReadTheDocs](https://readthedocs.io) using [this guide](https://docs.readthedocs.io/en/stable/intro/import-guide.html). The [`.readthedocs.yml`](.readthedocs.yml) YAML file contains all the configuration you'll need. You can also set up continuous integration on GitHub to check not only that Sphinx can build the documentation in an isolated environment (i.e., with `tox -e docs-test`) but also that [ReadTheDocs can build it too](https://docs.readthedocs.io/en/stable/pull-requests.html). #### Configuring ReadTheDocs 1. Log in to ReadTheDocs with your GitHub account to install the integration at https://readthedocs.org/accounts/login/?next=/dashboard/ 2. Import your project by navigating to https://readthedocs.org/dashboard/import then clicking the plus icon next to your repository 3. You can rename the repository on the next screen using a more stylized name (i.e., with spaces and capital letters) 4. Click next, and you're good to go! ### 📦 Making a Release #### Configuring Zenodo [Zenodo](https://zenodo.org) is a long-term archival system that assigns a DOI to each release of your package. 1. Log in to Zenodo via GitHub with this link: https://zenodo.org/oauth/login/github/?next=%2F. This brings you to a page that lists all of your organizations and asks you to approve installing the Zenodo app on GitHub. Click "grant" next to any organizations you want to enable the integration for, then click the big green "approve" button. This step only needs to be done once. 2. Navigate to https://zenodo.org/account/settings/github/, which lists all of your GitHub repositories (both in your username and any organizations you enabled). Click the on/off toggle for any relevant repositories. When you make a new repository, you'll have to come back to this After these steps, you're ready to go! After you make "release" on GitHub (steps for this are below), you can navigate to https://zenodo.org/account/settings/github/repository/cthoyt/pystow to see the DOI for the release and link to the Zenodo record for it. #### Registering with the Python Package Index (PyPI) You only have to do the following steps once. 1. Register for an account on the [Python Package Index (PyPI)](https://pypi.org/account/register) 2. Navigate to https://pypi.org/manage/account and make sure you have verified your email address. A verification email might not have been sent by default, so you might have to click the "options" dropdown next to your address to get to the "re-send verification email" button 3. 2-Factor authentication is required for PyPI since the end of 2023 (see this [blog post from PyPI](https://blog.pypi.org/posts/2023-05-25-securing-pypi-with-2fa/)). This means you have to first issue account recovery codes, then set up 2-factor authentication 4. Issue an API token from https://pypi.org/manage/account/token #### Configuring your machine's connection to PyPI You have to do the following steps once per machine. ```console $ uv tool install keyring $ keyring set https://upload.pypi.org/legacy/ __token__ $ keyring set https://test.pypi.org/legacy/ __token__ ``` Note that this deprecates previous workflows using `.pypirc`. #### Uploading to PyPI After installing the package in development mode and installing `tox` with `uv tool install tox --with tox-uv` or `python3 -m pip install tox tox-uv`, run the following from the console: ```console $ tox -e finish ``` This script does the following: 1. Uses [bump-my-version](https://github.com/callowayproject/bump-my-version) to switch the version number in the `pyproject.toml`, `CITATION.cff`, `src/pystow/version.py`, and [`docs/source/conf.py`](docs/source/conf.py) to not have the `-dev` suffix 2. Packages the code in both a tar archive and a wheel using [`uv build`](https://docs.astral.sh/uv/guides/publish/#building-your-package) 3. Uploads to PyPI using [`uv publish`](https://docs.astral.sh/uv/guides/publish/#publishing-your-package). 4. Push to GitHub. You'll need to make a release going with the commit where the version was bumped. 5. Bump the version to the next patch. If you made big changes and want to bump the version by minor, you can use `tox -e bumpversion -- minor` after. #### Releasing on GitHub 1. Navigate to https://github.com/cthoyt/pystow/releases/new to draft a new release 2. Click the "Choose a Tag" dropdown and select the tag corresponding to the release you just made 3. Click the "Generate Release Notes" button to get a quick outline of recent changes. Modify the title and description as you see fit 4. Click the big green "Publish Release" button This will trigger Zenodo to assign a DOI to your release as well.
cthoyt-pystow-54a586e/docs/000077500000000000000000000000001525577235500157025ustar00rootroot00000000000000cthoyt-pystow-54a586e/docs/Makefile000066400000000000000000000011351525577235500173420ustar00rootroot00000000000000# Minimal makefile for Sphinx documentation # You can set these variables from the command line. SPHINXOPTS = SPHINXBUILD = sphinx-build SPHINXPROJ = PyStow SOURCEDIR = source BUILDDIR = build # Put it first so that "make" without argument is like "make help". help: @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) .PHONY: help Makefile # Catch-all target: route all unknown targets to Sphinx using the new # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). %: Makefile @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)cthoyt-pystow-54a586e/docs/source/000077500000000000000000000000001525577235500172025ustar00rootroot00000000000000cthoyt-pystow-54a586e/docs/source/cli.rst000066400000000000000000000003521525577235500205030ustar00rootroot00000000000000######################## Command Line Interface ######################## pystow automatically installs the command ``pystow``. See ``pystow --help`` for usage details. .. click:: pystow.cli:main :prog: pystow :show-nested: cthoyt-pystow-54a586e/docs/source/conf.py000066400000000000000000000174651525577235500205160ustar00rootroot00000000000000"""Configuration file for the Sphinx documentation builder. This file does only contain a selection of the most common options. For a full list see the documentation: http://www.sphinx-doc.org/en/master/config -- Path setup -------------------------------------------------------------- If extensions (or modules to document with autodoc) are in another directory, add these directories to ``sys.path`` here. If the directory is relative to the documentation root, use ``os.path.abspath`` to make it absolute, like shown here. """ import os import re import sys from datetime import date sys.path.insert(0, os.path.abspath("../../src")) # -- Project information ----------------------------------------------------- project = "pystow" copyright = f"{date.today().year}, Charles Tapley Hoyt" # noqa:DTZ011 author = "Charles Tapley Hoyt" # The full version, including alpha/beta/rc tags. release = "0.9.8" # The short X.Y version. parsed_version = re.match( r"(?P\d+)\.(?P\d+)\.(?P\d+)(?:-(?P[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))?(?:\+(?P[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))?", release, ) version = parsed_version.expand(r"\g.\g.\g") if parsed_version.group("release"): tags.add("prerelease") # noqa:F821 # See https://about.readthedocs.com/blog/2024/07/addons-by-default/ # Define the canonical URL if you are using a custom domain on Read the Docs html_baseurl = os.environ.get("READTHEDOCS_CANONICAL_URL", "") # See https://about.readthedocs.com/blog/2024/07/addons-by-default/ # Tell Jinja2 templates the build is running on Read the Docs if os.environ.get("READTHEDOCS", "") == "True": if "html_context" not in globals(): html_context = {} html_context["READTHEDOCS"] = True # -- General configuration --------------------------------------------------- # If your documentation needs a minimal Sphinx version, state it here. # # needs_sphinx = '1.0' # If true, the current module name will be prepended to all description # unit titles (such as .. function::). add_module_names = False # A list of prefixes that are ignored when creating the module index. (new in Sphinx 0.6) modindex_common_prefix = ["pystow."] # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ "sphinx.ext.autosummary", "sphinx.ext.autodoc", "sphinx.ext.coverage", "sphinx.ext.intersphinx", "sphinx.ext.todo", "sphinx.ext.mathjax", "sphinx.ext.viewcode", "sphinx_automodapi.automodapi", "sphinx_automodapi.smart_resolver", ] extensions.append("sphinx_click.ext") # generate autosummary pages autosummary_generate = True # Add any paths that contain templates here, relative to this directory. templates_path = ["_templates"] # The suffix(es) of source filenames. # You can specify multiple suffix as a list of string: # # source_suffix = ['.rst', '.md'] source_suffix = { ".rst": "restructuredtext", } # The master toctree document. master_doc = "index" # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. # # This is also used if you do content translation via gettext catalogs. # Usually you set "language" from the command line for these cases. language = "en" # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This pattern also affects html_static_path and html_extra_path. exclude_patterns = [] # The name of the Pygments (syntax highlighting) style to use. pygments_style = "sphinx" # -- Options for HTML output ------------------------------------------------- # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. # html_theme = "sphinx_rtd_theme" # Theme options are theme-specific and customize the look and feel of a theme # further. For a list of options available for each theme, see the # documentation. # # html_theme_options = {} # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". # html_static_path = ['_static'] # Custom sidebar templates, must be a dictionary that maps document names # to template names. # # The default sidebars (for documents that don't match any pattern) are # defined by theme itself. Builtin themes are using these templates by # default: ``['localtoc.html', 'relations.html', 'sourcelink.html', # 'searchbox.html']``. # # html_sidebars = {} # The name of an image file (relative to this directory) to place at the top # of the sidebar. # if os.path.exists("logo.png"): html_logo = "logo.png" # -- Options for HTMLHelp output --------------------------------------------- # Output file base name for HTML help builder. htmlhelp_basename = "pystow_doc" # -- Options for LaTeX output ------------------------------------------------ # latex_elements = { # The paper size ('letterpaper' or 'a4paper'). # # 'papersize': 'letterpaper', # # The font size ('10pt', '11pt' or '12pt'). # # 'pointsize': '10pt', # # Additional stuff for the LaTeX preamble. # # 'preamble': '', # # Latex figure (float) alignment # # 'figure_align': 'htbp', # } # Grouping the document tree into LaTeX files. List of tuples # (source start file, target name, title, # author, documentclass [howto, manual, or own class]). # latex_documents = [ # ( # master_doc, # 'pystow.tex', # 'PyStow Documentation', # author, # 'manual', # ), # ] # -- Options for manual page output ------------------------------------------ # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). man_pages = [ ( master_doc, "pystow", "PyStow Documentation", [author], 1, ), ] # -- Options for Texinfo output ---------------------------------------------- # Grouping the document tree into Texinfo files. List of tuples # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ ( master_doc, "pystow", "PyStow Documentation", author, "Charles Tapley Hoyt", "Easily pick a place to store data for your Python code", "Miscellaneous", ), ] # -- Options for Epub output ------------------------------------------------- # Bibliographic Dublin Core info. # epub_title = project # The unique identifier of the text. This can be a ISBN number # or the project homepage. # # epub_identifier = '' # A unique identification for the text. # # epub_uid = '' # A list of files that should not be packed into the epub file. # epub_exclude_files = ['search.html'] # -- Extension configuration ------------------------------------------------- # -- Options for intersphinx extension --------------------------------------- # Example configuration for intersphinx: refer to the Python standard library. # Note: don't add trailing slashes, since sphinx adds "/objects.inv" to the end intersphinx_mapping = { "python": ("https://docs.python.org/3", None), "rdflib": ("https://rdflib.readthedocs.io/en/latest", None), "pandas": ("https://pandas.pydata.org/docs", None), "sklearn": ("https://scikit-learn.org/stable", None), "numpy": ("https://numpy.org/doc/stable", None), "scipy": ("https://docs.scipy.org/doc/scipy", None), } autoclass_content = "both" # Don't sort alphabetically, explained at: # https://stackoverflow.com/questions/37209921/python-how-not-to-sort-sphinx-output-in-alphabetical-order autodoc_member_order = "bysource" todo_include_todos = True todo_emit_warnings = True # Output SVG inheritance diagrams graphviz_output_format = "svg" cthoyt-pystow-54a586e/docs/source/index.rst000066400000000000000000000044271525577235500210520ustar00rootroot00000000000000################################ PyStow |release| Documentation ################################ If you've ever written the following few lines of code, :mod:`pystow` is for you: .. code-block:: python import os home = os.path.expanduser("~") project_name = "adeft" envvar_name = f"{project_name.upper()}_HOME" if envvar_name in os.environ: ADEFT_HOME = os.environ[envvar_name] else: ADEFT_HOME = os.path.join(home, f".{project_name}") os.makedirs(ADEFT_HOME, exist_ok=True) Many projects (let's use `Adeft `_ as an example) create a folder in the home directory as a dot-file such as ``$HOME/.adeft``. I found that I had so many of these that I started grouping them inside a ``$HOME/.data`` folder. It's also the case that every time you create one of these folders, you need to ensure its existence. :mod:`pystow` takes care of these things. You can replace the previous code with: .. code-block:: python import pystow ADEFT_HOME = pystow.join("adeft") First, it takes the name of the module, uppercases it, and postpends ``_HOME`` on to it (e.g., ``ADEFT_HOME``) and looks in the environment. If this variable is available, it uses that as the directory. It ensures it exists, then returns a :class:`pathlib.Path` pointing to it. If ``ADEFT_HOME`` (or more generally, ``_HOME`` is not available in the environment, it picks the path as ``$HOME/.data/``. Normally, ``$HOME`` is specified in your OS. However, if you want to pick another location to stick the data, you can override using ``$HOME`` by setting ``$PYSTOW_HOME`` in the environment. If you want to go more directories deep inside the adeft default directory, you can just keep using more positional arguments (the same semantics as :func:`os.path.join`). These directories automatically get created as well. .. code-block:: python import pystow from pathlib import Path # already set somewhere __version__ = ... ADEFT_VERSION_HOME: Path = pystow.join("adeft", __version__) .. toctree:: :maxdepth: 2 :caption: Getting Started :name: start installation usage utils cli ******************** Indices and Tables ******************** - :ref:`genindex` - :ref:`modindex` - :ref:`search` cthoyt-pystow-54a586e/docs/source/installation.rst000066400000000000000000000063751525577235500224500ustar00rootroot00000000000000############## Installation ############## The most recent release can be installed from `PyPI `_ with uv: .. code-block:: console $ uv pip install pystow or with pip: .. code-block:: console $ python3 -m pip install pystow ********************* Installing from git ********************* The most recent code and data can be installed directly from GitHub with uv: .. code-block:: console $ uv pip install git+https://github.com/cthoyt/pystow.git or with pip: .. code-block:: console $ python3 -m pip install git+https://github.com/cthoyt/pystow.git **************************** Installing for development **************************** To install in development mode with uv: .. code-block:: console $ git clone git+https://github.com/cthoyt/pystow.git $ cd pystow $ uv pip install -e . or with pip: .. code-block:: console $ python3 -m pip install -e . ############### Configuration ############### By default, data is stored in the ``$HOME/.data`` directory. By default, the ```` app will create the ``$HOME/.data/`` folder. If you want to use an alternate folder name to ``.data`` inside the home directory, you can set the ``PYSTOW_NAME`` environment variable. For example, if you set ``PYSTOW_NAME=mydata``, then the following code for the ``pykeen`` app will create the ``$HOME/mydata/pykeen/`` directory: .. code-block:: python import os import pystow # Only for demonstration purposes. You should set environment # variables either with your .bashrc or in the command line REPL. os.environ["PYSTOW_NAME"] = "mydata" # Get a directory (as a pathlib.Path) for ~/mydata/pykeen pykeen_directory = pystow.join("pykeen") If you want to specify a completely custom directory that isn't relative to your home directory, you can set the ``PYSTOW_HOME`` environment variable. For example, if you set ``PYSTOW_HOME=/usr/local/``, then the following code for the ``pykeen`` app will create the ``/usr/local/pykeen/`` directory: .. code-block:: python import os import pystow # Only for demonstration purposes. You should set environment # variables either with your .bashrc or in the command line REPL. os.environ["PYSTOW_HOME"] = "/usr/local/" # Get a directory (as a pathlib.Path) for /usr/local/pykeen pykeen_directory = pystow.join("pykeen") .. warning:: If you set ``PYSTOW_HOME``, then ``PYSTOW_NAME`` is disregarded. ************************************* X Desktop Group (XDG) Compatibility ************************************* While PyStow's main goal is to make application data less opaque and less hidden, some users might want to use the `XDG specifications `_ for storing their app data. If you set the environment variable ``PYSTOW_USE_APPDIRS`` to ``true`` or ``True``, then the `appdirs `_ or `platformdirs `_ package will be used to choose the base directory based on the ``user data dir`` option. .. warning:: If you use this setting, make sure you first do ``pip install appdirs`` or ``pip install platformdirs`` .. note:: This can still be overridden by ``PYSTOW_HOME``. cthoyt-pystow-54a586e/docs/source/usage.rst000066400000000000000000000004101525577235500210330ustar00rootroot00000000000000####### Usage ####### .. automodapi:: pystow :no-inheritance-diagram: :no-heading: :headings: -- :skip: Module :no-main-docstr: .. automodapi:: pystow.impl :no-inheritance-diagram: :no-heading: :headings: -- :no-main-docstr: cthoyt-pystow-54a586e/docs/source/utils.rst000066400000000000000000000002011525577235500210650ustar00rootroot00000000000000########### Utilities ########### .. automodapi:: pystow.utils :no-inheritance-diagram: :no-heading: :headings: -- cthoyt-pystow-54a586e/pyproject.toml000066400000000000000000000135311525577235500176710ustar00rootroot00000000000000[build-system] requires = ["uv_build>=0.6.6,<1.0"] build-backend = "uv_build" [project] name = "pystow" version = "0.9.8" description = "Easily pick a place to store data for your Python code" readme = "README.md" authors = [ { name = "Charles Tapley Hoyt", email = "cthoyt@gmail.com" } ] maintainers = [ { name = "Charles Tapley Hoyt", email = "cthoyt@gmail.com" } ] # See https://packaging.python.org/en/latest/guides/writing-pyproject-toml/#classifiers # Search tags using the controlled vocabulary at https://pypi.org/classifiers classifiers = [ "Development Status :: 5 - Production/Stable", "Environment :: Console", "Intended Audience :: Developers", "Operating System :: OS Independent", "Framework :: Pytest", "Framework :: tox", "Framework :: Sphinx", "Natural Language :: English", "Programming Language :: Python", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", "Programming Language :: Python :: 3 :: Only", "Typing :: Typed", ] keywords = [ "snekpack", # please keep this keyword to credit the cookiecutter-snekpack template "cookiecutter", "caching", "file management" ] # License Information. # See PEP-639 at https://peps.python.org/pep-0639/#add-license-files-key license = "MIT" license-files = ["LICENSE"] requires-python = ">=3.11" dependencies = [ "tqdm", "typing-extensions", "backports.zstd; python_version < '3.14'", ] [dependency-groups] tests = [ "pytest", "coverage[toml]", "requests-file", ] typing = [ "mypy", "pydantic", "types-requests", "types-ratelimit", "types-PyYAML", "types-lxml", ] docs = [ "sphinx>=8", "sphinx-rtd-theme>=3.0", "sphinx-click", "sphinx-automodapi", ] [project.optional-dependencies] rdf = ["rdflib"] xml = ["lxml"] pandas = ["pandas"] aws = ["boto3"] ratelimit = ["ratelimit", "requests"] bs4 = ["bs4", "requests"] yaml = ["pyyaml"] pydantic = ["pydantic"] cli = ["click"] requests = ["requests"] # See https://packaging.python.org/en/latest/guides/writing-pyproject-toml/#urls # and also https://packaging.python.org/en/latest/specifications/well-known-project-urls/ [project.urls] "Bug Tracker" = "https://github.com/cthoyt/pystow/issues" Homepage = "https://github.com/cthoyt/pystow" Repository = "https://github.com/cthoyt/pystow.git" Documentation = "https://pystow.readthedocs.io" Funding = "https://github.com/sponsors/cthoyt" [project.scripts] pystow = "pystow.cli:main" [tool.cruft] skip = [ "**/__init__.py", "tests/*" ] # MyPy, see https://mypy.readthedocs.io/en/stable/config_file.html [tool.mypy] plugins = [ ] # Doc8, see https://doc8.readthedocs.io/en/stable/readme.html#ini-file-usage [tool.doc8] max-line-length = 120 # Pytest, see https://docs.pytest.org/en/stable/reference/customize.html#pyproject-toml [tool.pytest.ini_options] markers = [ "slow: marks tests as slow (deselect with '-m \"not slow\"')", ] filterwarnings = [ "ignore::pyparsing.exceptions.PyparsingDeprecationWarning" ] # Coverage, see https://coverage.readthedocs.io/en/latest/config.html [tool.coverage.run] branch = true source = [ "pystow", ] omit = [ "tests/*", "docs/*", ] [tool.coverage.paths] source = [ "src/pystow", ".tox/*/lib/python*/site-packages/pystow", ] [tool.coverage.report] show_missing = true exclude_lines = [ "pragma: no cover", "raise NotImplementedError", "if __name__ == \"__main__\":", "if TYPE_CHECKING:", "def __str__", "def __repr__", ] [tool.ruff] line-length = 100 extend-include = ["*.ipynb"] [tool.ruff.lint] # See https://docs.astral.sh/ruff/rules extend-select = [ "F", # pyflakes "E", # pycodestyle errors "W", # pycodestyle warnings "C90", # mccabe "I", # isort "UP", # pyupgrade "D", # pydocstyle "DOC", # pydoclint "B", # bugbear "S", # bandit "T20", # print "N", # pep8 naming "ERA", # eradicate commented out code "NPY", # numpy checks "RUF", # ruff rules "C4", # comprehensions ] ignore = [ "D105", # Missing docstring in magic method "E203", # Black conflicts with the following "S301", # yolo pickle ] # See https://docs.astral.sh/ruff/settings/#per-file-ignores [tool.ruff.lint.per-file-ignores] # Ignore security issues in the version.py, which are inconsistent "src/pystow/version.py" = ["S603", "S607"] # Ignore commented out code in Sphinx configuration file "docs/source/conf.py" = ["ERA001"] # Prints are okay in notebooks "notebooks/**/*.ipynb" = ["T201"] [tool.ruff.lint.pydocstyle] convention = "pep257" [tool.ruff.lint.isort] relative-imports-order = "closest-to-furthest" known-third-party = [ "tqdm", ] known-first-party = [ "pystow", "tests", ] [tool.ruff.format] # see https://docs.astral.sh/ruff/settings/#format_docstring-code-format docstring-code-format = true [tool.bumpversion] current_version = "0.9.8" parse = "(?P\\d+)\\.(?P\\d+)\\.(?P\\d+)(?:-(?P[0-9A-Za-z-]+(?:\\.[0-9A-Za-z-]+)*))?(?:\\+(?P[0-9A-Za-z-]+(?:\\.[0-9A-Za-z-]+)*))?" serialize = [ "{major}.{minor}.{patch}-{release}+{build}", "{major}.{minor}.{patch}+{build}", "{major}.{minor}.{patch}-{release}", "{major}.{minor}.{patch}", ] commit = true tag = false [tool.bumpversion.parts.release] optional_value = "production" first_value = "dev" values = [ "dev", "production", ] [[tool.bumpversion.files]] filename = "pyproject.toml" search = "version = \"{current_version}\"" replace = "version = \"{new_version}\"" [[tool.bumpversion.files]] filename = "docs/source/conf.py" search = "release = \"{current_version}\"" replace = "release = \"{new_version}\"" [[tool.bumpversion.files]] filename = "src/pystow/version.py" search = "VERSION = \"{current_version}\"" replace = "VERSION = \"{new_version}\"" cthoyt-pystow-54a586e/src/000077500000000000000000000000001525577235500155415ustar00rootroot00000000000000cthoyt-pystow-54a586e/src/pystow/000077500000000000000000000000001525577235500171065ustar00rootroot00000000000000cthoyt-pystow-54a586e/src/pystow/__init__.py000066400000000000000000000042131525577235500212170ustar00rootroot00000000000000"""PyStow: Easily pick a place to store data for your python package.""" from .api import ( dump_df, dump_json, dump_pickle, dump_rdf, dump_xml, ensure, ensure_csv, ensure_custom, ensure_excel, ensure_from_google, ensure_from_s3, ensure_gunzip, ensure_json, ensure_json_bz2, ensure_nltk, ensure_open, ensure_open_bz2, ensure_open_gz, ensure_open_lzma, ensure_open_sqlite, ensure_open_sqlite_gz, ensure_open_tarfile, ensure_open_zip, ensure_open_zstd, ensure_pickle, ensure_pickle_gz, ensure_rdf, ensure_soup, ensure_tar_df, ensure_tar_xml, ensure_untar, ensure_xml, ensure_yaml, ensure_zip_df, ensure_zip_np, get_sentence_transformer, join, joinpath_sqlite, load_df, load_json, load_pickle, load_pickle_gz, load_rdf, load_xml, load_yaml, module, open, open_gz, ) from .config_api import ConfigError, get_config, write_config from .impl import Module, VersionHint from .utils import ensure_readme __all__ = [ "ConfigError", "Module", "VersionHint", "dump_df", "dump_json", "dump_pickle", "dump_rdf", "dump_xml", "ensure", "ensure_csv", "ensure_custom", "ensure_excel", "ensure_from_google", "ensure_from_s3", "ensure_gunzip", "ensure_json", "ensure_json_bz2", "ensure_nltk", "ensure_open", "ensure_open_bz2", "ensure_open_gz", "ensure_open_lzma", "ensure_open_sqlite", "ensure_open_sqlite_gz", "ensure_open_tarfile", "ensure_open_zip", "ensure_open_zstd", "ensure_pickle", "ensure_pickle_gz", "ensure_rdf", "ensure_soup", "ensure_tar_df", "ensure_tar_xml", "ensure_untar", "ensure_xml", "ensure_yaml", "ensure_zip_df", "ensure_zip_np", "get_config", "get_sentence_transformer", "join", "joinpath_sqlite", "load_df", "load_json", "load_pickle", "load_pickle_gz", "load_rdf", "load_xml", "load_yaml", "module", "open", "open_gz", "write_config", ] ensure_readme() del ensure_readme cthoyt-pystow-54a586e/src/pystow/__main__.py000066400000000000000000000001671525577235500212040ustar00rootroot00000000000000# type: ignore """Command line interface for PyStow.""" from .cli import main if __name__ == "__main__": main() cthoyt-pystow-54a586e/src/pystow/api.py000066400000000000000000002144051525577235500202370ustar00rootroot00000000000000"""API functions for PyStow.""" from __future__ import annotations import bz2 import io import lzma import sqlite3 import typing from collections.abc import Generator, Mapping, Sequence from contextlib import contextmanager from functools import lru_cache from io import BytesIO, StringIO from pathlib import Path from typing import IO, TYPE_CHECKING, Any, Literal, overload from .constants import JSON, Provider from .impl import Module, VersionHint from .utils.download import DownloadKwargs if TYPE_CHECKING: from pathlib import PurePath import bs4 import lxml.etree import numpy.typing import pandas as pd import rdflib import sentence_transformers __all__ = [ "DEFAULT_SENTENCE_TRANSFORMER_NAME", "dump_df", "dump_json", "dump_pickle", "dump_rdf", "dump_xml", "ensure", "ensure_csv", "ensure_custom", "ensure_excel", "ensure_from_google", "ensure_from_s3", "ensure_gunzip", "ensure_json", "ensure_json_bz2", "ensure_nltk", "ensure_open", "ensure_open_bz2", "ensure_open_gz", "ensure_open_lzma", "ensure_open_sqlite", "ensure_open_sqlite_gz", "ensure_open_tarfile", "ensure_open_zip", "ensure_open_zstd", "ensure_pickle", "ensure_pickle_gz", "ensure_rdf", "ensure_soup", "ensure_tar_df", "ensure_tar_xml", "ensure_untar", "ensure_xml", "ensure_yaml", "ensure_zip_df", "ensure_zip_np", "get_sentence_transformer", "join", "joinpath_sqlite", "load_df", "load_json", "load_pickle", "load_pickle_gz", "load_rdf", "load_xml", "load_yaml", "module", "open", "open_gz", ] def module(key: str, *subkeys: str, ensure_exists: bool = True) -> Module: """Return a module for the application. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param ensure_exists: Should all directories be created automatically? Defaults to true. :returns: The module object that manages getting and ensuring """ return Module.from_key(key, *subkeys, ensure_exists=ensure_exists) def join( key: str, *subkeys: str, name: str | None = None, ensure_exists: bool = True, version: VersionHint = None, ) -> Path: """Return the home data directory for the given module. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join :param name: The name of the file (optional) inside the folder :param ensure_exists: Should all directories be created automatically? Defaults to true. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. The following example describes how to store the versioned data from the Rhea database for biologically relevant chemical reactions. .. code-block:: import pystow import requests def get_rhea_version() -> str: res = requests.get("https://ftp.expasy.org/databases/rhea/rhea-release.properties") _, _, version = res.text.splitlines()[0].partition("=") return version # Assume you want to download the data from # ftp://ftp.expasy.org/databases/rhea/rdf/rhea.rdf.gz, make a path # with the same name path = pystow.join("rhea", name="rhea.rdf.gz", version=get_rhea_version) :returns: The path of the directory or subdirectory for the given module. """ _module = Module.from_key(key, ensure_exists=ensure_exists) return _module.join(*subkeys, name=name, ensure_exists=ensure_exists, version=version) # docstr-coverage:excused `overload` @overload @contextmanager def open( key: str, *subkeys: str, name: str, mode: Literal["r", "rt", "w", "wt"] = "r", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[StringIO]: ... # docstr-coverage:excused `overload` @overload @contextmanager def open( key: str, *subkeys: str, name: str, mode: Literal["rb", "wb"], open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[BytesIO]: ... @contextmanager def open( key: str, *subkeys: str, name: str, mode: Literal["r", "rb", "rt", "w", "wb", "wt"] = "r", open_kwargs: Mapping[str, Any] | None = None, ensure_exists: bool = False, ) -> Generator[StringIO | BytesIO]: """Open a file. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param mode: The read or write mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param ensure_exists: Should the directory the file is in be made? Set to true on write operations. :yields: An open file object This function should be called inside a context manager like in the following .. code-block:: python import pystow with pystow.open("test", name="test.tsv", mode="w") as file: print("Test text!", file=file) """ _module = Module.from_key(key, ensure_exists=True) with _module.open( *subkeys, name=name, mode=mode, open_kwargs=open_kwargs, ensure_exists=ensure_exists ) as file: yield file # docstr-coverage:excused `overload` @overload @contextmanager def open_gz( key: str, *subkeys: str, name: str, mode: Literal["r", "w", "rt", "wt"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[StringIO]: ... # docstr-coverage:excused `overload` @overload @contextmanager def open_gz( key: str, *subkeys: str, name: str, mode: Literal["rb", "wb"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[BytesIO]: ... @contextmanager def open_gz( key: str, *subkeys: str, name: str, mode: Literal["r", "w", "rt", "wt", "rb", "wb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, ensure_exists: bool = False, ) -> Generator[StringIO | BytesIO]: """Open a gzipped file that exists already. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param mode: The read mode, passed to :func:`gzip.open` :param open_kwargs: Additional keyword arguments passed to :func:`gzip.open` :param ensure_exists: Should the file be made? Set to true on write operations. :yields: An open file object """ _module = Module.from_key(key, ensure_exists=True) with _module.open_gz( *subkeys, name=name, mode=mode, open_kwargs=open_kwargs, ensure_exists=ensure_exists ) as file: yield file def ensure( key: str, *subkeys: str, url: str, name: str | None = None, version: VersionHint = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, ) -> Path: """Ensure a file is downloaded. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. The following example describes how to store the versioned data from the Rhea database for biologically relevant chemical reactions. .. code-block:: import pystow import requests def get_rhea_version() -> str: res = requests.get("https://ftp.expasy.org/databases/rhea/rhea-release.properties") _, _, version = res.text.splitlines()[0].partition("=") return version path = pystow.ensure( "rhea", url="ftp://ftp.expasy.org/databases/rhea/rdf/rhea.rdf.gz", version=get_rhea_version, ) :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :returns: The path of the file that has been downloaded (or already exists) """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure( *subkeys, url=url, name=name, version=version, force=force, download_kwargs=download_kwargs ) def ensure_soup( key: str, *subkeys: str, url: str, name: str | None = None, version: VersionHint = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, beautiful_soup_kwargs: Mapping[str, Any] | None = None, ) -> bs4.BeautifulSoup: """Ensure a webpage is downloaded and parsed with :mod:`BeautifulSoup`. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param beautiful_soup_kwargs: Additional keyword arguments passed to :class:`BeautifulSoup` :returns: An BeautifulSoup object .. note:: If you don't need to cache, consider using :func:`pystow.utils.get_soup` instead. """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_soup( *subkeys, url=url, name=name, version=version, force=force, download_kwargs=download_kwargs, beautiful_soup_kwargs=beautiful_soup_kwargs, ) def ensure_custom( key: str, *subkeys: str, name: str, force: bool = False, provider: Provider, **kwargs: Any, ) -> Path: """Ensure a file is present, and run a custom create function otherwise. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The file name. :param force: Should the file be re-created, even if the path already exists? :param provider: The file provider. Will be run with the path as the first positional argument, if the file needs to be generated. :param kwargs: Additional keyword-based parameters passed to the provider. :returns: The path of the file that has been created (or already exists) """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_custom(*subkeys, name=name, force=force, provider=provider, **kwargs) def ensure_untar( key: str, *subkeys: str, url: str, name: str | None = None, directory: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, extract_kwargs: Mapping[str, Any] | None = None, ) -> Path: """Ensure a file is downloaded and untarred. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param directory: Overrides the name of the directory into which the tar archive is extracted. If none given, will use the stem of the file name that gets downloaded. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param extract_kwargs: Keyword arguments to pass to :meth:`tarfile.TarFile.extract_all`. :returns: The path of the directory where the file that has been downloaded gets extracted to """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_untar( *subkeys, url=url, name=name, directory=directory, force=force, download_kwargs=download_kwargs, extract_kwargs=extract_kwargs, ) def ensure_gunzip( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, autoclean: bool = True, download_kwargs: DownloadKwargs | None = None, ) -> Path: """Ensure a file is downloaded and gunzipped. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param autoclean: Should the zipped file be deleted? :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :returns: The path of the directory where the file that has been downloaded gets extracted to """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_gunzip( *subkeys, url=url, name=name, force=force, autoclean=autoclean, download_kwargs=download_kwargs, ) # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open( key: str, *subkeys: str, url: str, name: str | None, force: bool, download_kwargs: DownloadKwargs | None, mode: Literal["r", "rt", "w", "wt"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[StringIO]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open( key: str, *subkeys: str, url: str, name: str | None, force: bool, download_kwargs: DownloadKwargs | None, mode: Literal["rb", "wb"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[BytesIO]: ... @contextmanager def ensure_open( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rt", "w", "wt", "rb", "wb"] = "r", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[StringIO | BytesIO]: """Ensure a file is downloaded and open it. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`lzma.open` :param open_kwargs: Additional keyword arguments passed to :func:`lzma.open` :yields: An open file object """ _module = Module.from_key(key, ensure_exists=True) with _module.ensure_open( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, ) as yv: yield yv # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_zip( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = ..., force: bool = ..., download_kwargs: DownloadKwargs | None = ..., mode: Literal["r", "w", "rb", "wb"] = ..., zipfile_kwargs: Mapping[str, Any] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., ) -> Generator[IO[bytes]]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_zip( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = ..., force: bool = ..., download_kwargs: DownloadKwargs | None = ..., mode: Literal["rt", "wt"] = ..., zipfile_kwargs: Mapping[str, Any] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., ) -> Generator[IO[str]]: ... @contextmanager def ensure_open_zip( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "w", "rb", "wb", "rt", "wt"] = "r", zipfile_kwargs: Mapping[str, Any] | None = None, open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[IO[str]] | Generator[IO[bytes]]: """Ensure a file is downloaded then open it with :mod:`zipfile`. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`zipfile.open` :param zipfile_kwargs: Additional keyword arguments passed to :class:`zipfile.ZipFile` :param open_kwargs: Additional keyword arguments passed to :func:`zipfile.open` :yields: An open file object """ _module = Module.from_key(key, ensure_exists=True) # ignore the call overload because mypy doesn't understand # when multiple overloads cascade (i.e., this function's # overloads correspond to the _module.ensure_open_zip's # overloads with _module.ensure_open_zip( # type:ignore[misc,call-overload] *subkeys, url=url, inner_path=inner_path, name=name, force=force, download_kwargs=download_kwargs, mode=mode, zipfile_kwargs=zipfile_kwargs, open_kwargs=open_kwargs, ) as yv: yield yv # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_lzma( key: str, *subkeys: str, url: str, name: str | None, force: bool, download_kwargs: DownloadKwargs | None, mode: Literal["r", "w", "rt", "wt"] = "rt", open_kwargs: Mapping[str, Any] | None, ) -> Generator[io.TextIOWrapper[lzma.LZMAFile]]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_lzma( key: str, *subkeys: str, url: str, name: str | None, force: bool, download_kwargs: DownloadKwargs | None, mode: Literal["rb", "wb"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[lzma.LZMAFile]: ... @contextmanager def ensure_open_lzma( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rb", "w", "wb", "rt", "wt"] = "rt", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[lzma.LZMAFile | io.TextIOWrapper[lzma.LZMAFile]]: """Ensure a LZMA-compressed file is downloaded and open a file inside it. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`lzma.open` :param open_kwargs: Additional keyword arguments passed to :func:`lzma.open` :yields: An open file object """ _module = Module.from_key(key, ensure_exists=True) with _module.ensure_open_lzma( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, ) as yv: yield yv @contextmanager def ensure_open_zstd( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rb", "w", "wb", "rt", "wt"] = "rt", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[io.BufferedIOBase]: """Ensure a zstd-compressed file is downloaded and open a file inside it. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`zstd.open` :param open_kwargs: Additional keyword arguments passed to :func:`zstd.open` :yields: An open file object """ _module = Module.from_key(key, ensure_exists=True) with _module.ensure_open_zstd( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, ) as yv: yield yv # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_tarfile( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = ..., force: bool = False, download_kwargs: DownloadKwargs | None = ..., mode: Literal["rt"] = ..., open_kwargs: Mapping[str, Any] | None = ..., ) -> Generator[IO[str]]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_tarfile( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = ..., force: bool = False, download_kwargs: DownloadKwargs | None = ..., mode: Literal["r", "rb"] = ..., open_kwargs: Mapping[str, Any] | None = ..., ) -> Generator[typing.IO[bytes]]: ... @contextmanager def ensure_open_tarfile( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rb", "rt"] = "r", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[IO[str]] | Generator[typing.IO[bytes]]: """Ensure a tar file is downloaded and open a file inside it. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`tarfile.open` :param open_kwargs: Additional keyword arguments passed to :func:`tarfile.open` :yields: An open file object """ _module = Module.from_key(key, ensure_exists=True) with _module.ensure_open_tarfile( *subkeys, url=url, inner_path=inner_path, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, ) as yv: yield yv # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_gz( key: str, *subkeys: str, url: str, name: str | None, force: bool = False, download_kwargs: DownloadKwargs | None, mode: Literal["r", "w", "rt", "wt"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[StringIO]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_gz( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["rb", "wb"] = ..., open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[BytesIO]: ... @contextmanager def ensure_open_gz( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rb", "w", "wb", "rt", "wt"] = "rb", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[StringIO | BytesIO]: """Ensure a gzipped file is downloaded and open a file inside it. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`gzip.open` :param open_kwargs: Additional keyword arguments passed to :func:`gzip.open` :yields: An open file object """ _module = Module.from_key(key, ensure_exists=True) with _module.ensure_open_gz( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, ) as yv: yield yv @contextmanager def ensure_open_bz2( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[bz2.BZ2File]: """Ensure a BZ2-compressed file is downloaded and open a file inside it. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`bz2.open` :param open_kwargs: Additional keyword arguments passed to :func:`bz2.open` :yields: An open file object """ _module = Module.from_key(key, ensure_exists=True) with _module.ensure_open_bz2( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, ) as yv: yield yv def ensure_csv( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, read_csv_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Download a CSV and open as a dataframe with :mod:`pandas`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param read_csv_kwargs: Keyword arguments to pass through to :func:`pandas.read_csv`. .. note:: It is assumed that the CSV uses tab separators, as this is the only safe option. For more information, see `Wikipedia `_ and `Issue #51 `_. To override this behavior and load using the comma separator, specify ``read_csv_kwargs=dict(sep=",")``. :returns: A pandas DataFrame Example usage: :: >>> import pystow >>> import pandas as pd >>> url = "https://raw.githubusercontent.com/pykeen/pykeen/master/src/pykeen/datasets/nations/test.txt" >>> df: pd.DataFrame = pystow.ensure_csv("pykeen", "datasets", "nations", url=url) """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_csv( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, read_csv_kwargs=read_csv_kwargs, ) def load_df( key: str, *subkeys: str, name: str, read_csv_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Open a pre-existing CSV as a dataframe with :mod:`pandas`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param read_csv_kwargs: Keyword arguments to pass through to :func:`pandas.read_csv`. :returns: A pandas DataFrame Example usage: :: >>> import pystow >>> import pandas as pd >>> url = "https://raw.githubusercontent.com/pykeen/pykeen/master/src/pykeen/datasets/nations/test.txt" >>> pystow.ensure_csv("pykeen", "datasets", "nations", url=url) >>> df: pd.DataFrame = pystow.load_df("pykeen", "datasets", "nations", name="test.txt") """ _module = Module.from_key(key, ensure_exists=True) return _module.load_df( *subkeys, name=name, read_csv_kwargs=read_csv_kwargs, ) def ensure_yaml( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, open_kwargs: Mapping[str, Any] | None = None, yaml_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Download YAML and open with :mod:`yaml`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param open_kwargs: Additional keyword arguments passed to :func:`open` :param yaml_load_kwargs: Keyword arguments to pass through to :func:`yaml.safe_load`. :returns: A JSON object (list, dict, etc.) """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_yaml( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, open_kwargs=open_kwargs, yaml_load_kwargs=yaml_load_kwargs, ) def dump_df( key: str, *subkeys: str, name: str, obj: pd.DataFrame, sep: str = "\t", index: bool = False, to_csv_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump a dataframe to a TSV file with :mod:`pandas`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param obj: The dataframe to dump :param sep: The separator to use, defaults to a tab :param index: Should the index be dumped? Defaults to false. :param to_csv_kwargs: Keyword arguments to pass through to :meth:`pandas.DataFrame.to_csv`. """ _module = Module.from_key(key, ensure_exists=True) _module.dump_df( *subkeys, name=name, obj=obj, sep=sep, index=index, to_csv_kwargs=to_csv_kwargs, ) def ensure_json( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, version: VersionHint = None, download_kwargs: DownloadKwargs | None = None, open_kwargs: Mapping[str, Any] | None = None, json_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Download JSON and open with :mod:`json`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param open_kwargs: Additional keyword arguments passed to :func:`open` :param json_load_kwargs: Keyword arguments to pass through to :func:`json.load`. :returns: A JSON object (list, dict, etc.) Example usage: :: >>> import pystow >>> url = "https://maayanlab.cloud/CREEDS/download/single_gene_perturbations-v1.0.json" >>> perturbations = pystow.ensure_json("bio", "creeds", "1.0", url=url) """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_json( *subkeys, url=url, name=name, force=force, version=version, download_kwargs=download_kwargs, open_kwargs=open_kwargs, json_load_kwargs=json_load_kwargs, ) def ensure_json_bz2( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, version: VersionHint = None, download_kwargs: DownloadKwargs | None = None, open_kwargs: Mapping[str, Any] | None = None, json_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Download BZ2-compressed JSON and open with :mod:`json`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param open_kwargs: Additional keyword arguments passed to :func:`bz2.open` :param json_load_kwargs: Keyword arguments to pass through to :func:`json.load`. :returns: A JSON object (list, dict, etc.) Example usage: :: >>> import pystow >>> url = "https://github.com/hetio/hetionet/raw/master/hetnet/json/hetionet-v1.0.json.bz2" >>> hetionet = pystow.ensure_json_bz2("bio", "hetionet", "1.0", url=url) """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_json_bz2( *subkeys, url=url, name=name, force=force, version=version, download_kwargs=download_kwargs, open_kwargs=open_kwargs, json_load_kwargs=json_load_kwargs, ) def load_yaml( key: str, *subkeys: str, name: str, yaml_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Open a JSON file :mod:`json`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param yaml_load_kwargs: Keyword arguments to pass through to :func:`yaml.safe_load`. :returns: A JSON object (list, dict, etc.) """ _module = Module.from_key(key, ensure_exists=True) return _module.load_yaml(*subkeys, name=name, yaml_load_kwargs=yaml_load_kwargs) def load_json( key: str, *subkeys: str, name: str, json_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Open a JSON file :mod:`json`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param json_load_kwargs: Keyword arguments to pass through to :func:`json.load`. :returns: A JSON object (list, dict, etc.) """ _module = Module.from_key(key, ensure_exists=True) return _module.load_json(*subkeys, name=name, json_load_kwargs=json_load_kwargs) def dump_json( key: str, *subkeys: str, name: str, obj: JSON, open_kwargs: Mapping[str, Any] | None = None, json_dump_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump an object to a file with :mod:`json`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param obj: The object to dump :param open_kwargs: Additional keyword arguments passed to :func:`open` :param json_dump_kwargs: Keyword arguments to pass through to :func:`json.dump`. """ _module = Module.from_key(key, ensure_exists=True) _module.dump_json( *subkeys, name=name, obj=obj, open_kwargs=open_kwargs, json_dump_kwargs=json_dump_kwargs ) def ensure_pickle( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, pickle_load_kwargs: Mapping[str, Any] | None = None, ) -> Any: """Download a pickle file and open with :mod:`pickle`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param pickle_load_kwargs: Keyword arguments to pass through to :func:`pickle.load`. :returns: Any object """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_pickle( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, pickle_load_kwargs=pickle_load_kwargs, ) def load_pickle( key: str, *subkeys: str, name: str, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, pickle_load_kwargs: Mapping[str, Any] | None = None, ) -> Any: """Open a pickle file with :mod:`pickle`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param pickle_load_kwargs: Keyword arguments to pass through to :func:`pickle.load`. :returns: Any object """ _module = Module.from_key(key, ensure_exists=True) return _module.load_pickle( *subkeys, name=name, mode=mode, open_kwargs=open_kwargs, pickle_load_kwargs=pickle_load_kwargs, ) def dump_pickle( key: str, *subkeys: str, name: str, obj: Any, mode: Literal["wb"] = "wb", open_kwargs: Mapping[str, Any] | None = None, pickle_dump_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump an object to a file with :mod:`pickle`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param obj: The object to dump :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param pickle_dump_kwargs: Keyword arguments to pass through to :func:`pickle.dump`. """ _module = Module.from_key(key, ensure_exists=True) _module.dump_pickle( *subkeys, name=name, obj=obj, mode=mode, open_kwargs=open_kwargs, pickle_dump_kwargs=pickle_dump_kwargs, ) def ensure_pickle_gz( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, pickle_load_kwargs: Mapping[str, Any] | None = None, ) -> Any: """Download a gzipped pickle file and open with :mod:`pickle`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`gzip.open` :param open_kwargs: Additional keyword arguments passed to :func:`gzip.open` :param pickle_load_kwargs: Keyword arguments to pass through to :func:`pickle.load`. :returns: Any object """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_pickle_gz( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, pickle_load_kwargs=pickle_load_kwargs, ) def load_pickle_gz( key: str, *subkeys: str, name: str, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, pickle_load_kwargs: Mapping[str, Any] | None = None, ) -> Any: """Open a gzipped pickle file with :mod:`pickle`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`gzip.open` :param pickle_load_kwargs: Keyword arguments to pass through to :func:`pickle.load`. :returns: Any object """ _module = Module.from_key(key, ensure_exists=True) return _module.load_pickle_gz( *subkeys, name=name, mode=mode, open_kwargs=open_kwargs, pickle_load_kwargs=pickle_load_kwargs, ) def ensure_xml( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, parse_kwargs: Mapping[str, Any] | None = None, ) -> lxml.etree.ElementTree: """Download an XML file and open it with :mod:`lxml`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param parse_kwargs: Keyword arguments to pass through to :func:`lxml.etree.parse`. :returns: An ElementTree object .. warning:: If you have lots of files to read in the same archive, it's better just to unzip first. """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_xml( *subkeys, name=name, url=url, force=force, download_kwargs=download_kwargs, parse_kwargs=parse_kwargs, ) def load_xml( key: str, *subkeys: str, name: str, parse_kwargs: Mapping[str, Any] | None = None, ) -> lxml.etree.ElementTree: """Load an XML file with :mod:`lxml`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param parse_kwargs: Keyword arguments to pass through to :func:`lxml.etree.parse`. :returns: An ElementTree object .. warning:: If you have lots of files to read in the same archive, it's better just to unzip first. """ _module = Module.from_key(key, ensure_exists=True) return _module.load_xml( *subkeys, name=name, parse_kwargs=parse_kwargs, ) def dump_xml( key: str, *subkeys: str, name: str, obj: lxml.etree.ElementTree, open_kwargs: Mapping[str, Any] | None = None, write_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump an XML element tree to a file with :mod:`lxml`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param obj: The object to dump :param open_kwargs: Additional keyword arguments passed to :func:`open` :param write_kwargs: Keyword arguments to pass through to :func:`lxml.etree.ElementTree.write`. """ _module = Module.from_key(key, ensure_exists=True) _module.dump_xml( *subkeys, name=name, obj=obj, open_kwargs=open_kwargs, write_kwargs=write_kwargs, ) def ensure_excel( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, read_excel_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Download an Excel file and open as a dataframe with :mod:`pandas`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param read_excel_kwargs: Keyword arguments to pass through to :func:`pandas.read_excel`. :returns: A pandas DataFrame """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_excel( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, read_excel_kwargs=read_excel_kwargs, ) def ensure_tar_df( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, read_csv_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Download a tar file and open an inner file as a dataframe with :mod:`pandas`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param read_csv_kwargs: Keyword arguments to pass through to :func:`pandas.read_csv`. :returns: A dataframe .. warning:: If you have lots of files to read in the same archive, it's better just to unzip first. """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_tar_df( *subkeys, url=url, name=name, force=force, inner_path=inner_path, download_kwargs=download_kwargs, read_csv_kwargs=read_csv_kwargs, ) def ensure_tar_xml( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, parse_kwargs: Mapping[str, Any] | None = None, ) -> lxml.etree.ElementTree: """Download a tar file and open an inner file as an XML with :mod:`lxml`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param parse_kwargs: Keyword arguments to pass through to :func:`lxml.etree.parse`. :returns: An ElementTree object .. warning:: If you have lots of files to read in the same archive, it's better just to unzip first. """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_tar_xml( *subkeys, url=url, name=name, force=force, inner_path=inner_path, download_kwargs=download_kwargs, parse_kwargs=parse_kwargs, ) def ensure_zip_df( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, read_csv_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Download a zip file and open an inner file as a dataframe with :mod:`pandas`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param read_csv_kwargs: Keyword arguments to pass through to :func:`pandas.read_csv`. :returns: A pandas DataFrame """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_zip_df( *subkeys, url=url, name=name, force=force, inner_path=inner_path, download_kwargs=download_kwargs, read_csv_kwargs=read_csv_kwargs, ) def ensure_zip_np( key: str, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, load_kwargs: Mapping[str, Any] | None = None, ) -> numpy.typing.ArrayLike: """Download a zip file and open an inner file as an array-like with :mod:`numpy`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param load_kwargs: Additional keyword arguments that are passed through to :func:`read_zip_np` and transitively to :func:`numpy.load`. :returns: An array-like object """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_zip_np( *subkeys, url=url, name=name, force=force, inner_path=inner_path, download_kwargs=download_kwargs, load_kwargs=load_kwargs, ) def ensure_rdf( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, precache: bool = True, parse_kwargs: Mapping[str, Any] | None = None, ) -> rdflib.Graph: """Download a RDF file and open with :mod:`rdflib`. :param key: The module name :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param precache: Should the parsed :class:`rdflib.Graph` be stored as a pickle for fast loading? :param parse_kwargs: Keyword arguments to pass through to :func:`pystow.utils.read_rdflib` and transitively to :func:`rdflib.Graph.parse`. :returns: An RDF graph Example usage .. code-block:: python import pystow import rdflib url = "https://ftp.expasy.org/databases/rhea/rdf/rhea.rdf.gz" rdf_graph: rdflib.Graph = pystow.ensure_rdf("rhea", url=url, parse_kwargs={"format": "xml"}) .. note:: Sometimes, :mod:`rdflib` is able to guess the format, and you can omit the "format" from the `parse_kwargs` argument. Here's another example .. code-block:: python import pystow import rdflib url = "http://oaei.webdatacommons.org/tdrs/testdata/persistent/knowledgegraph/v3/suite/memoryalpha-stexpanded/component/reference.xml" rdf_graph: rdflib.Graph = pystow.ensure_rdf( "memoryalpha-stexpanded", url=url, parse_kwargs={"format": "xml"} ) """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_rdf( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, precache=precache, parse_kwargs=parse_kwargs, ) def load_rdf( key: str, *subkeys: str, name: str | None = None, parse_kwargs: Mapping[str, Any] | None = None, ) -> rdflib.Graph: """Open an RDF file with :mod:`rdflib`. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param parse_kwargs: Keyword arguments to pass through to :func:`pystow.utils.read_rdflib` and transitively to :func:`rdflib.Graph.parse`. :returns: An RDF graph """ _module = Module.from_key(key, ensure_exists=True) return _module.load_rdf(*subkeys, name=name, parse_kwargs=parse_kwargs) def dump_rdf( key: str, *subkeys: str, name: str, obj: rdflib.Graph, format: str = "turtle", serialize_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump an RDF graph to a file with :mod:`rdflib`. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param obj: The object to dump :param format: The format to dump in :param serialize_kwargs: Keyword arguments to through to :func:`rdflib.Graph.serialize`. """ _module = Module.from_key(key, ensure_exists=True) _module.dump_rdf(*subkeys, name=name, obj=obj, format=format, serialize_kwargs=serialize_kwargs) def ensure_from_s3( key: str, *subkeys: str, s3_bucket: str, s3_key: str | Sequence[str], name: str | None = None, force: bool = False, **kwargs: Any, ) -> Path: """Ensure a file is downloaded. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param s3_bucket: The S3 bucket name :param s3_key: The S3 key name :param name: Overrides the name of the file at the end of the S3 key, if given. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param kwargs: Remaining kwargs to forwrad to :class:`Module.ensure_from_s3`. :returns: The path of the file that has been downloaded (or already exists) Example downloading ProtMapper 0.0.21: .. code-block:: python import pystow version = "0.0.21" pystow.ensure_from_s3( "test", version, s3_bucket="bigmech", s3_key=f"protmapper/{version}/refseq_uniprot.csv", ) """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_from_s3( *subkeys, s3_bucket=s3_bucket, s3_key=s3_key, name=name, force=force, **kwargs ) def ensure_from_google( key: str, *subkeys: str, name: str, file_id: str, force: bool = False, ) -> Path: """Ensure a file is downloaded from google drive. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file :param file_id: The file identifier of the google file. If your share link is https://drive.google.com/file/d/1AsPPU4ka1Rc9u-XYMGWtvV65hF3egi0z/view, then your file id is ``1AsPPU4ka1Rc9u-XYMGWtvV65hF3egi0z``. :param force: Should the download be done again, even if the path already exists? Defaults to false. :returns: The path of the file that has been downloaded (or already exists) Example downloading the WK3l-15k dataset as motivated by https://github.com/pykeen/pykeen/pull/403: .. code-block:: python import pystow path = pystow.ensure_from_google( "test", name="wk3l15k.zip", file_id="1AsPPU4ka1Rc9u-XYMGWtvV65hF3egi0z" ) """ _module = Module.from_key(key, ensure_exists=True) return _module.ensure_from_google(*subkeys, name=name, file_id=file_id, force=force) def joinpath_sqlite(key: str, *subkeys: str, name: str) -> str: """Get an SQLite database connection string. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the database file. :returns: A SQLite path string. """ _module = Module.from_key(key, ensure_exists=True) return _module.joinpath_sqlite(*subkeys, name=name) @contextmanager def ensure_open_sqlite( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, ) -> Generator[sqlite3.Connection]: """Ensure and connect to a SQLite database. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :yields: An instance of :class:`sqlite3.Connection` from :func:`sqlite3.connect` Example usage: .. code-block:: python import pystow import pandas as pd url = "https://s3.amazonaws.com/bbop-sqlite/obi.db" sql = "SELECT * FROM entailed_edge LIMIT 10" with pystow.ensure_open_sqlite("test", url=url) as conn: df = pd.read_sql(sql, conn) """ _module = Module.from_key(key, ensure_exists=True) with _module.ensure_open_sqlite( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) as yv: yield yv @contextmanager def ensure_open_sqlite_gz( key: str, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, ) -> Generator[sqlite3.Connection]: """Ensure and connect to a gzipped SQLite database. :param key: The name of the module. No funny characters. The envvar `_HOME` where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :yields: An instance of :class:`sqlite3.Connection` from :func:`sqlite3.connect` Example usage: >>> import pystow >>> import pandas as pd >>> url = "https://s3.amazonaws.com/bbop-sqlite/hp.db.gz" >>> sql = "SELECT * FROM entailed_edge LIMIT 10" >>> with pystow.ensure_open_sqlite_gz("test", url=url) as conn: >>> df = pd.read_sql(sql, conn) """ _module = Module.from_key(key, ensure_exists=True) with _module.ensure_open_sqlite_gz( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) as yv: yield yv @lru_cache def ensure_nltk(resource: str = "stopwords") -> tuple[Path, bool]: """Ensure NLTK data is downloaded in a standard way. :param resource: Name of the resource to download, e.g., ``stopwords`` :returns: A pair of the NLTK cache directory and a boolean that says if download was successful This function also appends the standard PyStow location for NLTK data to the :data:`nltk.data.path` list so any downstream users of NLTK will know how to find it automatically. """ import nltk.data directory = join("nltk") result = nltk.download(resource, download_dir=directory, quiet=True) if directory not in nltk.data.path: nltk.data.path.append(directory) # this is cached so you don't have to keep checking # if the package was downloaded return directory, result #: The default sentence transformer model DEFAULT_SENTENCE_TRANSFORMER_NAME = "all-MiniLM-L6-v2" def get_sentence_transformer( name: str | sentence_transformers.SentenceTransformer | None = None, **kwargs: Any ) -> sentence_transformers.SentenceTransformer: """Get a sentence transformer. :param name: The name of the sentence transformer model on HuggingFace. If none is passed, then :data:`DEFAULT_SENTENCE_TRANSFORMER_NAME` is used. If a pre-instantiated model is passed, then it is passed through. :param kwargs: Keyword arguments to pass to :class:`sentence_transformers.SentenceTransformer`. :returns: An instantiated sentence transformer object, which has a :meth:`sentence_transformers.SentenceTransformer.encode` function """ from sentence_transformers import SentenceTransformer if name is None: name = DEFAULT_SENTENCE_TRANSFORMER_NAME elif not isinstance(name, str): # allow a pre-loaded model to go through return name directory = join("sentence-transformers", name) model = SentenceTransformer(name, cache_folder=directory.as_posix(), **kwargs) return model cthoyt-pystow-54a586e/src/pystow/cache.py000066400000000000000000000153511525577235500205300ustar00rootroot00000000000000"""Utilities for caching files.""" from __future__ import annotations import functools import json import logging import pickle from abc import ABC, abstractmethod from collections.abc import Callable, MutableMapping from pathlib import Path from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast from .utils import read_pydantic_json, write_pydantic_json if TYPE_CHECKING: import pandas as pd import pydantic __all__ = [ # Classes "Cached", "CachedCollection", "CachedDataFrame", "CachedJSON", "CachedPickle", "CachedPydantic", # Types "Getter", ] logger = logging.getLogger(__name__) JSONType = dict[str, Any] | list[Any] X = TypeVar("X") Getter = Callable[[], X] class Cached(ABC, Generic[X]): """Caching decorator.""" def __init__( self, path: str | Path, *, force: bool = False, cache: bool = True, ) -> None: """Instantiate the decorator. :param path: The path to the cache for the file :param cache: Should caching be done? Defaults to true, turn off for debugging purposes :param force: Should a pre-existing file be disregared/overwritten? """ self.path = Path(path) self.force = force self.cache = cache def __call__(self, func: Getter[X]) -> Getter[X]: """Apply this instance as a decorator. :param func: The function to wrap :returns: A wrapped function """ @functools.wraps(func) def _wrapped() -> X: if not self.cache: return func() if self.path.is_file() and not self.force: return self.load() logger.debug("no cache found at %s", self.path) rv = func() logger.debug("writing cache to %s", self.path) self.dump(rv) return rv return _wrapped @abstractmethod def load(self) -> X: """Load data from the cache (typically by opening a file at the given path).""" @abstractmethod def dump(self, rv: X) -> None: """Dump data to the cache (typically by opening a file at the given path). :param rv: The data to dump """ class CachedJSON(Cached[JSONType]): """Make a function lazily cache its return value as JSON.""" def load(self) -> JSONType: """Load data from the cache as JSON. :returns: A python object with JSON-like data from the cache """ with open(self.path) as file: return cast(JSONType, json.load(file)) def dump(self, rv: JSONType) -> None: """Dump data to the cache as JSON. :param rv: The JSON data to dump """ with open(self.path, "w") as file: json.dump(rv, file, indent=2) B = TypeVar("B", bound="pydantic.BaseModel") class CachedPydantic(Cached[B]): """Make a function lazily cache its return value as JSON.""" model_cls: type[B] def __init__( self, path: str | Path, *, model_cls: type[B], force: bool = False, cache: bool = True, ) -> None: """Instantiate the decorator.""" super().__init__(path=path, force=force, cache=cache) self.model_cls = model_cls def load(self) -> B: """Load data from the cache as JSON. :returns: A Pydantic model parsed from the JSON """ return read_pydantic_json(self.path, self.model_cls) def dump(self, model: B) -> None: """Dump data to the cache as JSON. :param model: The Pydantic model to dump """ write_pydantic_json(model, self.path) class CachedPickle(Cached[Any]): """Make a function lazily cache its return value as a pickle.""" def load(self) -> Any: """Load data from the cache as a pickle. :returns: A python object loaded from the cache """ with open(self.path, "rb") as file: return pickle.load(file) def dump(self, rv: Any) -> None: """Dump data to the cache as a pickle. :param rv: The arbitrary python object to dump """ with open(self.path, "wb") as file: pickle.dump(rv, file, protocol=pickle.HIGHEST_PROTOCOL) class CachedCollection(Cached[list[str]]): """Make a function lazily cache its return value as file.""" def load(self) -> list[str]: """Load data from the cache as a list of strings. :returns: A list of strings loaded from the cache """ with open(self.path) as file: return [line.strip() for line in file] def dump(self, rv: list[str]) -> None: """Dump data to the cache as a list of strings. :param rv: The list of strings to dump """ with open(self.path, "w") as file: for line in rv: print(line, file=file) class CachedDataFrame(Cached["pd.DataFrame"]): """Make a function lazily cache its return value as a dataframe.""" def __init__( self, path: str | Path, cache: bool = True, force: bool = False, sep: str | None = None, dtype: Any | None = None, read_csv_kwargs: MutableMapping[str, Any] | None = None, ) -> None: """Instantiate the decorator. :param path: The path to the cache for the file :param force: Should a pre-existing file be disregared/overwritten? :param sep: The separator. Defaults to TSV, since this is the only reasonable default. :param dtype: A shortcut for setting the dtype :param read_csv_kwargs: Additional kwargs to pass to :func:`pd.read_csv`. :raises ValueError: if sep is given as a kwarg and also in ``read_csv_kwargs``. """ super().__init__(path=path, cache=cache, force=force) self.read_csv_kwargs = read_csv_kwargs or {} if "sep" not in self.read_csv_kwargs: self.sep = sep or "\t" elif sep is not None: raise ValueError else: self.sep = self.read_csv_kwargs.pop("sep") if dtype is not None: if "dtype" in self.read_csv_kwargs: raise ValueError self.read_csv_kwargs["dtype"] = dtype self.read_csv_kwargs.setdefault("keep_default_na", False) def load(self) -> pd.DataFrame: """Load data from the cache as a dataframe. :returns: A dataframe loaded from the cache. """ import pandas as pd return pd.read_csv( self.path, sep=self.sep, **self.read_csv_kwargs, ) def dump(self, rv: pd.DataFrame) -> None: """Dump data to the cache as a dataframe. :param rv: The dataframe to dump """ rv.to_csv(self.path, sep=self.sep, index=False) cthoyt-pystow-54a586e/src/pystow/cli.py000066400000000000000000000027531525577235500202360ustar00rootroot00000000000000"""Command line interface for PyStow.""" from __future__ import annotations import os from collections.abc import Sequence from pathlib import Path import click @click.group() def main() -> None: """Run the PyStow CLI.""" @main.command() @click.argument("keys", nargs=-1) @click.option("--name") def join(keys: Sequence[str], name: str | None) -> None: """List a directory.""" from . import api click.echo(api.join(*keys, name=name)) @main.command() @click.argument("keys", nargs=-1) def ls(keys: Sequence[str]) -> None: """List a directory.""" from . import api directory = api.join(*keys) _ls(directory) @main.command() @click.argument("keys", nargs=-1) @click.option("--url", required=True) @click.option("--name") @click.option("--force", is_flag=True) def ensure(keys: Sequence[str], url: str, name: str | None, force: bool) -> None: """Ensure a file is downloaded.""" from . import api path = api.ensure(*keys, url=url, name=name, force=force) _ls(path.parent) def _ls(directory: Path) -> None: command = f"ls -al {directory}" click.secho(f"[pystow] {command}", fg="cyan", bold=True) os.system(command) # noqa:S605 @main.command(name="set") @click.argument("module") @click.argument("key") @click.argument("value") def set_config(module: str, key: str, value: str) -> None: """Set a configuration value.""" from .config_api import write_config write_config(module, key, value) if __name__ == "__main__": main() cthoyt-pystow-54a586e/src/pystow/config_api.py000066400000000000000000000160131525577235500215570ustar00rootroot00000000000000"""Configuration handling.""" from __future__ import annotations import os from collections.abc import Callable from configparser import ConfigParser from functools import lru_cache from pathlib import Path from textwrap import dedent from typing import Any, TypeVar from .utils import getenv_path __all__ = [ "get_config", "write_config", ] X = TypeVar("X") CONFIG_NAME_ENVVAR = "PYSTOW_CONFIG_NAME" CONFIG_HOME_ENVVAR = "PYSTOW_CONFIG_HOME" CONFIG_NAME_DEFAULT = ".config" class ConfigError(ValueError): """Raised when configuration can not be looked up.""" def __init__(self, module: str, key: str): """Initialize the configuration error. :param module: Name of the module, e.g., ``bioportal`` :param key: Name of the key inside the module, e.g., ``api_key`` """ self.module = module self.key = key def __str__(self) -> str: path = get_home().joinpath(self.module).with_suffix(".ini") return dedent( f"""\ Could not look up {self.module}/{self.key} and no default given. This can be solved with one of the following: 1. Set the {self.module.upper()}_{self.key.upper()} environment variable - Windows, via GUI: https://www.computerhope.com/issues/ch000549.htm - Windows, via CLI: https://learn.microsoft.com/en-us/windows-server/administration/windows-commands/set_1 - Mac OS: https://apple.stackexchange.com/questions/106778/how-do-i-set-environment-variables-on-os-x - Linux: https://www.freecodecamp.org/news/how-to-set-an-environment-variable-in-linux/ 2. Use the PyStow CLI from the command line to set the configuration like so: $ pystow set {self.module} {self.key} This creates an INI file in {path} with the configuration in the right place. 3. Create/edit an INI file in {path} and manually fill it in by 1) creating a section inside it called [{self.module}] and 2) setting a value for {self.key} = that looks like: # {path} [{self.module}] {self.key} = See https://github.com/cthoyt/pystow#%EF%B8%8F%EF%B8%8F-configuration for more information. """ ) def get_name() -> str: """Get the config home directory name. :returns: The name of the pystow home directory, either loaded from the :data:`CONFIG_NAME_ENVVAR` environment variable or given by the default value :data:`CONFIG_NAME_DEFAULT`. """ return os.getenv(CONFIG_NAME_ENVVAR, default=CONFIG_NAME_DEFAULT) def get_home(ensure_exists: bool = True) -> Path: """Get the config home directory. :param ensure_exists: If true, ensures the directory is created :returns: A path object representing the pystow home directory, as one of: 1. :data:`CONFIG_HOME_ENVVAR` environment variable or 2. The default directory constructed in the user's home directory plus what's returned by :func:`get_name`. """ default = Path.home().joinpath(get_name()).expanduser() return getenv_path(CONFIG_HOME_ENVVAR, default, ensure_exists=ensure_exists) @lru_cache(maxsize=1) def _get_cfp(module: str) -> ConfigParser: cfp = ConfigParser() directory = get_home() # If a multi-part module was given like "zenodo:sandbox", # then only look for the first part "zenodo" as the file name if ":" in module: module = module.split(":", 1)[0] filenames = [ os.path.join(directory, "config.cfg"), os.path.join(directory, "config.ini"), os.path.join(directory, "pystow.cfg"), os.path.join(directory, "pystow.ini"), os.path.join(directory, f"{module}.cfg"), os.path.join(directory, f"{module}.ini"), os.path.join(directory, module, f"{module}.cfg"), os.path.join(directory, module, f"{module}.ini"), os.path.join(directory, module, "conf.ini"), os.path.join(directory, module, "config.ini"), os.path.join(directory, module, "conf.cfg"), os.path.join(directory, module, "config.cfg"), ] cfp.read(filenames) return cfp def get_config( module: str, key: str, *, passthrough: X | None = None, default: X | None = None, dtype: type[X] | None = None, raise_on_missing: bool = False, ) -> Any: """Get a configuration value. :param module: Name of the module (e.g., ``pybel``) to get configuration for :param key: Name of the key (e.g., ``connection``) :param passthrough: If this is not none, will get returned :param default: If the environment and configuration files don't contain anything, this is returned. :param dtype: The datatype to parse out. Can either be :func:`int`, :func:`float`, :func:`bool`, or :func:`str`. If none, defaults to :func:`str`. :param raise_on_missing: If true, will raise a value error if no data is found and no default is given :returns: The config value or the default. :raises ConfigError: If ``raise_on_missing`` conditions are met """ if passthrough is not None: return _cast(passthrough, dtype) rv = os.getenv(f"{module.upper()}_{key.upper()}") if rv is not None: return _cast(rv, dtype) rv = _get_cfp(module).get(module, key, fallback=None) if rv is None: if default is None and raise_on_missing: raise ConfigError(module=module, key=key) return default return _cast(rv, dtype) def _cast(rv: Any, dtype: Callable[..., Any] | None) -> Any: if not isinstance(rv, str): # if it's not a string, it doesn't need munging return rv if dtype in (None, str): # no munging necessary return rv if dtype in (int, float): return dtype(rv) if dtype is bool: if rv.lower() in ("t", "true", "yes", "1", 1, True): return True elif rv.lower() in ("f", "false", "no", "0", 0, False): return False else: raise ValueError(f"value can not be coerced into bool: {rv}") raise TypeError(f"dtype is invalid: {dtype}") def write_config(module: str, key: str, value: str) -> None: """Write a configuration value. :param module: The name of the app (e.g., ``indra``) :param key: The key of the configuration in the app :param value: The value of the configuration in the app """ _get_cfp.cache_clear() cfp = ConfigParser() # If there's a multi-part module such as "zenodo:sandbox", # then write to zenodo.ini with section [zenodo:sandbox] fname = module.split(":", 1)[0] if ":" in module else module path = get_home().joinpath(fname).with_suffix(".ini") cfp.read(path) # If the file did not exist, then this section will be empty # and running set() would raise a configparser.NoSectionError. if not cfp.has_section(module): cfp.add_section(module) # Note that the section duplicates the file name cfp.set(section=module, option=key, value=value) with path.open("w") as file: cfp.write(file) cthoyt-pystow-54a586e/src/pystow/constants.py000066400000000000000000000043231525577235500214760ustar00rootroot00000000000000"""PyStow constants.""" from __future__ import annotations from collections.abc import Callable from textwrap import dedent from typing import Any, TypeAlias __all__ = [ "JSON", "PYSTOW_HOME_ENVVAR", "PYSTOW_NAME_DEFAULT", "PYSTOW_NAME_ENVVAR", "PYSTOW_USE_APPDIRS", "README_TEXT", "Provider", "TimeoutHint", ] PYSTOW_NAME_ENVVAR = "PYSTOW_NAME" PYSTOW_HOME_ENVVAR = "PYSTOW_HOME" PYSTOW_USE_APPDIRS = "PYSTOW_USE_APPDIRS" PYSTOW_NAME_DEFAULT = ".data" README_TEXT = dedent( """\ # PyStow Data Directory This directory is used by [`pystow`](https://github.com/cthoyt/pystow) as a reproducible location to store and access data. ### ⚙️️ Configuration By default, data is stored in the `$HOME/.data` directory. By default, the `` app will create the `$HOME/.data/` folder. If you want to use an alternate folder name to `.data` inside the home directory, you can set the `PYSTOW_NAME` environment variable. For example, if you set `PYSTOW_NAME=mydata`, then the following code for the `pykeen` app will create the `$HOME/mydata/pykeen/` directory: ```python import os import pystow # Only for demonstration purposes. You should set environment # variables either with your .bashrc or in the command line REPL. os.environ['PYSTOW_NAME'] = 'mydata' # Get a directory (as a pathlib.Path) for ~/mydata/pykeen pykeen_directory = pystow.join('pykeen') ``` If you want to specify a completely custom directory that isn't relative to your home directory, you can set the `PYSTOW_HOME` environment variable. For example, if you set `PYSTOW_HOME=/usr/local/`, then the following code for the `pykeen` app will create the `/usr/local/pykeen/` directory: ```python import os import pystow # Only for demonstration purposes. You should set environment # variables either with your .bashrc or in the command line REPL. os.environ['PYSTOW_HOME'] = '/usr/local/' # Get a directory (as a pathlib.Path) for /usr/local/pykeen pykeen_directory = pystow.join('pykeen') ``` Note: if you set `PYSTOW_HOME`, then `PYSTOW_NAME` is disregarded. """ ) JSON = Any Provider = Callable[..., None] #: A hint for timeout in :func:`requests.get` TimeoutHint: TypeAlias = int | float | tuple[float | int, float | int] | None cthoyt-pystow-54a586e/src/pystow/git.py000066400000000000000000000062611525577235500202500ustar00rootroot00000000000000"""Utilities for the web app.""" from __future__ import annotations from contextlib import contextmanager from typing import TYPE_CHECKING, TypeAlias if TYPE_CHECKING: from collections.abc import Generator from pathlib import Path from subprocess import CompletedProcess __all__ = [ "clone_github_tempdir", "clone_tempdir", "commit", "create_branch", "fetch", "get_current_branch", "git", "has_local_branch", "is_likely_default_branch", "push", ] @contextmanager def clone_github_tempdir(owner: str, repo: str) -> Generator[Path]: """Temporarily clone a repository from a URL.""" url = f"https://github.com/{owner}/{repo}.git" with clone_tempdir(url) as directory: yield directory @contextmanager def clone_tempdir(url: str) -> Generator[Path]: """Temporarily clone a repository from a URL.""" import shutil import tempfile from pathlib import Path with tempfile.TemporaryDirectory() as directory_: directory = Path(directory_) clone(directory, url) yield directory shutil.rmtree(directory) GitReturn: TypeAlias = "CompletedProcess[str]" def git(directory: Path, *args: str) -> GitReturn: """Run the git command with the given arguments in the given directory.""" return _check_output("git", *args, directory=directory) def clone(directory: Path, url: str) -> GitReturn: """Clone the git repository with the given URL.""" args = ["git", "clone", url, directory.as_posix()] return _check_output(*args, directory=directory) def _check_output(*args: str, directory: Path | None = None) -> GitReturn: from subprocess import run return run( # noqa: S603 args, cwd=None if directory is None else directory.as_posix(), capture_output=True, text=True, check=False, ) def commit(directory: Path, message: str) -> GitReturn: """Make a commit with the following message.""" return git(directory, "commit", "-m", message, "-a") def push(directory: Path, branch: str | None = None) -> GitReturn: """Push the git repo.""" if branch is not None: return git(directory, "push", "origin", branch) else: return git(directory, "push") def fetch(directory: Path) -> GitReturn: """Fetch the git repo.""" return git(directory, "fetch", "--all") def has_local_branch(directory: Path, branch: str) -> bool: """Check if a branch exists in the git repo.""" completed_process = git(directory, "show-ref", "--verify", "--quiet", f"refs/heads/{branch}") return completed_process.returncode == 0 def get_current_branch(directory: Path) -> str: """Return if on the master/main branch.""" completed_process = git(directory, "rev-parse", "--abbrev-ref", "HEAD") return completed_process.stdout.strip() def create_branch(directory: Path, branch: str) -> GitReturn: """Create a branch and check it out.""" return git(directory, "checkout", "-b", branch) LIKELY_DEFAULT_BRANCHES = {"master", "main"} def is_likely_default_branch(directory: Path) -> bool: """Guess if the current branch is the "default".""" return get_current_branch(directory) in LIKELY_DEFAULT_BRANCHES cthoyt-pystow-54a586e/src/pystow/github.py000066400000000000000000000176641525577235500207600ustar00rootroot00000000000000"""Interactions with GitHub.""" from __future__ import annotations from collections.abc import Iterable from functools import partial from typing import TYPE_CHECKING, Any, Literal, cast from ratelimit import rate_limited from tqdm import tqdm from .config_api import get_config from .constants import TimeoutHint if TYPE_CHECKING: import requests __all__ = [ "MAXIMUM_SEARCH_PAGE_SIZE", "delete_branch", "get_contents", "get_contributions", "get_default_branch", "get_issue", "get_issues", "get_pull_requests", "get_repository", "get_repository_commit_activity", "get_token", "get_topics", "get_user_events", "post_pull", "requests_get_github", "requests_post_github", "search_code", ] def get_token(*, passthrough: str | None = None, raise_on_missing: bool = False) -> str | None: """Get a GitHub token.""" rv = get_config("github", "token", passthrough=passthrough, raise_on_missing=raise_on_missing) return rv # type:ignore def get_headers( token: str | None = None, *, raise_on_missing: bool = False, preview: bool = False ) -> dict[str, str]: """Get GitHub headers.""" headers = { "Accept": "application/vnd.github.mercy-preview+json" if preview else "application/vnd.github.v3+json", } token = get_token(passthrough=token, raise_on_missing=raise_on_missing) if token: headers["Authorization"] = f"token {token}" return headers def requests_get_github(path: str, **kwargs: Any) -> requests.Response: """Make a GET request to the GitHub API.""" return _request_github("GET", path, **kwargs) def requests_post_github(path: str, **kwargs: Any) -> requests.Response: """Make a POST request to the GitHub API.""" return _request_github("POST", path, **kwargs) def requests_delete_github(path: str, **kwargs: Any) -> requests.Response: """Make a DELETE request to the GitHub API.""" return _request_github("DELETE", path, **kwargs) @rate_limited(calls=5_000, period=60 * 60) def _request_github( method: Literal["GET", "POST", "DELETE"], path: str, params: dict[str, Any] | None = None, token: str | None = None, timeout: TimeoutHint = None, require_token: bool = False, preview: bool = False, headers: dict[str, str] | None = None, **kwargs: Any, ) -> requests.Response: """Make a POST request to the GitHub API.""" import requests path = path.lstrip("/") url = f"https://api.github.com/{path}" headers_ = get_headers(token=token, raise_on_missing=require_token, preview=preview) if headers: headers_.update(headers) return requests.request(method, url, headers=headers_, params=params, timeout=timeout, **kwargs) def post_pull( owner: str, repo: str, *, title: str | None = None, head: str, head_repo: str | None = None, base: str, body: str | None = None, maintainer_can_modify: bool = True, draft: bool = False, issue: int | None = None, **kwargs: Any, ) -> requests.Response: """Post a pull request.""" # https://docs.github.com/en/rest/pulls/pulls?apiVersion=2022-11-28#create-a-pull-request if title is None and issue is None: raise ValueError("title and issue can't both be none") if "params" in kwargs: raise ValueError("should pass arguments directly") return requests_post_github( f"repos/{owner}/{repo}/pulls", params={ "title": title, "head": head, "head_repo": head_repo, "base": base, "body": body, "maintainer_can_modify": maintainer_can_modify, "draft": draft, "issue": issue, }, **kwargs, ) def delete_branch(owner: str, repo: str, branch_name: str, **kwargs: Any) -> requests.Response: """Delete a branch. :param owner: The repository owner :param repo: The repository name :param branch_name: The branch name or head ref :returns: A response from the GitHub API after running the (secret) delete endpoint """ # see https://github.com/orgs/community/discussions/24603 return requests_delete_github(f"repos/{owner}/{repo}/git/refs/heads/{branch_name}", **kwargs) def get_repository(owner: str, repo: str, **kwargs: Any) -> requests.Response: """Get information about a repository.""" return requests_get_github(f"repos/{owner}/{repo}", **kwargs) def get_default_branch(owner: str, repo: str, **kwargs: Any) -> str: """Get the default branch for the repository.""" res = get_repository(owner, repo, **kwargs).json() return cast(str, res["default_branch"]) def get_issues(owner: str, repo: str, **kwargs: Any) -> requests.Response: """Get issues from a repository.""" return requests_get_github(f"repos/{owner}/{repo}/issues", **kwargs) def get_issue(owner: str, repo: str, issue: int, **kwargs: Any) -> requests.Response: """Get an issue from a repository.""" return requests_get_github(f"repos/{owner}/{repo}/issues/{issue}", **kwargs) def get_pull_requests(owner: str, repo: str, **kwargs: Any) -> requests.Response: """Get pull requests from a repository.""" return requests_get_github(f"repos/{owner}/{repo}/pulls", **kwargs) def get_user_events(user: str) -> requests.Response: """Get events for a user.""" return requests_get_github(f"users/{user}/events") def get_contributions(owner: str, repo: str, **kwargs: Any) -> requests.Response: """Get contributors to a repository.""" return requests_get_github(f"repos/{owner}/{repo}/stats/contributors", **kwargs) def get_repository_commit_activity(owner: str, repo: str, **kwargs: Any) -> requests.Response: """Get commit activity to a repository.""" return requests_get_github(f"repos/{owner}/{repo}/stats/commit_activity", **kwargs) def get_topics(owner: str, repo: str, *, preview: bool = True, **kwargs: Any) -> requests.Response: """Get topics from the repository.""" return requests_get_github(f"repos/{owner}/{repo}/topics", preview=preview, **kwargs) #: Maximum number of records per page in code search MAXIMUM_SEARCH_PAGE_SIZE = 40 def search_code( query: str, *, page_size: int | None = None, progress: bool = True, inner_progress: bool = True, ) -> Iterable[dict[str, Any]]: """Search GitHub code.""" page = 1 if page_size is None: page_size = MAXIMUM_SEARCH_PAGE_SIZE page_size = min(page_size, MAXIMUM_SEARCH_PAGE_SIZE) inner_tqdm = partial(tqdm, disable=not inner_progress, unit="record", leave=False) initial_res = _search_code_helper(page_size=page_size, page=page, query=query) initial_res.raise_for_status() initial_res_json = initial_res.json() total = initial_res_json["total_count"] total_pages = 1 + (total // page_size) yield from inner_tqdm(initial_res_json["items"], desc="Page 1") with tqdm( total=total_pages, unit="page", disable=not progress, desc="Paginating code search results" ) as tbar: tbar.update(1) while page_size * page < total: page += 1 tbar.update(1) res = _search_code_helper(page=page, page_size=page_size, query=query) res.raise_for_status() res_json = res.json() yield from inner_tqdm(res_json["items"], desc=f"Page {page}") def _search_code_helper(page_size: int, page: int, query: str) -> requests.Response: return requests_get_github( "search/code", params={ "per_page": page_size, "page": page, "sort": "indexed", "q": query, }, ) def get_contents(owner: str, repo: str, path: str, **kwargs: Any) -> list[dict[str, Any]]: """List all files in a repository in the given path.""" res = requests_get_github( f"repos/{owner}/{repo}/contents/{path}", headers={"X-GitHub-Api-Version": "2026-03-10"}, **kwargs, ) res.raise_for_status() return cast(list[dict[str, Any]], res.json()) cthoyt-pystow-54a586e/src/pystow/graph.py000066400000000000000000000176621525577235500205750ustar00rootroot00000000000000"""Implementation of an efficient cached graph data structure.""" from __future__ import annotations from collections.abc import Callable, Collection, Iterable from dataclasses import dataclass from pathlib import Path from typing import Literal, Self, TypeAlias import numpy as np from tqdm import tqdm from .utils import safe_open __all__ = [ "GraphCache", "GraphCachePaths", "build_graph_cache", ] MissingAction: TypeAlias = Literal["raise", "none", "list"] @dataclass class GraphCachePaths: """An object with the paths for cached graph data.""" nodes: Path forward_indices_pointer: Path forward_indices: Path reverse_indices_pointer: Path reverse_indices: Path def exists(self) -> bool: """Check that the cache exists.""" return all( path.is_file() for path in ( self.nodes, self.forward_indices_pointer, self.reverse_indices_pointer, self.forward_indices, self.reverse_indices, ) ) @classmethod def from_directory(cls, directory: str | Path) -> Self: """Instantiate the object from the given directory.""" directory = Path(directory).expanduser().resolve() if not directory.is_dir(): raise NotADirectoryError return cls( nodes=directory / "nodes.txt.gz", forward_indices_pointer=directory / "fwd_indptr.bin", forward_indices=directory / "fwd_indices.bin", reverse_indices_pointer=directory / "rev_indptr.bin", reverse_indices=directory / "rev_indices.bin", ) class SingleGraphCache: """A tool for a one-directional graph.""" def __init__( self, index_pointer_path: Path, index_path: Path, node_to_id: dict[str, int], id_to_node: dict[int, str], ) -> None: """Construct a memory graph.""" self.node_to_id = node_to_id self.id_to_node = id_to_node self.indices_pointers = np.memmap(index_pointer_path, dtype=np.int64, mode="r") self.indices = np.memmap(index_path, dtype=np.int32, mode="r") def edges(self, node: str, *, raise_on_missing: bool = False) -> list[str]: """Get edges for the node.""" # is not None is important vs. regular falsy check, because one node will have id = 0 if (node_id := self.node_to_id.get(node)) is not None: return [self.id_to_node[neighbor_id] for neighbor_id in self._get_edges(node_id)] elif raise_on_missing: raise KeyError(f"{node} is not in the graph cache.") from None else: return [] def _get_edges(self, node: int) -> Collection[int]: return self.indices[self.indices_pointers[node] : self.indices_pointers[node + 1]] class GraphCache: """A tool for looking up in and out edges quickly.""" def __init__(self, paths: GraphCachePaths) -> None: """Construct a memory graph.""" with safe_open(paths.nodes) as file: node_to_id = {node.strip(): i for i, node in enumerate(file)} id_to_node = {v: k for k, v in node_to_id.items()} self.forward = SingleGraphCache( paths.forward_indices_pointer, paths.forward_indices, node_to_id=node_to_id, id_to_node=id_to_node, ) self.reverse = SingleGraphCache( paths.reverse_indices_pointer, paths.reverse_indices, node_to_id=node_to_id, id_to_node=id_to_node, ) @classmethod def from_directory(cls, directory: str | Path) -> Self: """Construct a memory graph from a directory.""" return cls(GraphCachePaths.from_directory(directory)) def in_edges(self, node: str, *, raise_on_missing: bool = False) -> list[str]: """Get in-edges for the node.""" return self.reverse.edges(node, raise_on_missing=raise_on_missing) def out_edges(self, node: str, *, raise_on_missing: bool = False) -> list[str]: """Get out-edges for the node.""" return self.forward.edges(node, raise_on_missing=raise_on_missing) def build_graph_cache( edges: Callable[[], Iterable[tuple[str, str]]], directory: str | Path | GraphCachePaths, *, sort_nodes: bool = False, progress: bool = True, estimated_edges: int | None = None, ) -> GraphCache: """Cache the directed graph to the disk using a CSR data structure. :param edges: A function, that when called, produces an iterable of edges (pairs of strings). This function accepts an iterable because it needs to make three passes over the edges, and to support large edge lists, this might take an iterator that reads a file. :param directory: The directory to use for caching, or a pre-instantiated cache directory object :param sort_nodes: Whether to sort the nodes first. This is not really needed in case of testing. :param progress: Whether to show a progress bar on the three passes of :param estimated_edges: On the first pass of the edge iterator, it's unknown how many edges there are. If you know, or have a good estimate, use this parameter to give more information to the progress bar. :returns: A graph cache object, which can access the written binaries quickly """ if not callable(edges): raise TypeError( "`edges` argument must be callable. This is because construction " "takes three passes, so it's better that a function that can iterate " "is given, to avoid needing to load into memory. If you already have " "your graph in memory, pass edges=lambda: edges`" ) if isinstance(directory, GraphCachePaths): paths = directory else: paths = GraphCachePaths.from_directory(directory) nodes: set[str] = set() n_edges = 0 for edge in tqdm( edges(), unit="edge", unit_scale=True, desc="indexing nodes", disable=not progress, total=estimated_edges, ): nodes.update(edge) n_edges += 1 n_nodes = len(nodes) tqdm.write(f"observed {n_nodes:,} nodes and {n_edges:,} edges") node_to_id = {} with safe_open(paths.nodes, operation="write") as file: for i, node in enumerate(sorted(nodes) if sort_nodes else nodes): node_to_id[node] = i print(node, file=file) forward_indices_pointer = np.zeros(n_nodes + 1, dtype=np.int64) reverse_indices_pointer = np.zeros(n_nodes + 1, dtype=np.int64) for u_str, v_str in tqdm( edges(), total=n_edges, unit="edge", unit_scale=True, desc="constructing pointers", disable=not progress, ): forward_indices_pointer[node_to_id[u_str] + 1] += 1 reverse_indices_pointer[node_to_id[v_str] + 1] += 1 np.cumsum(forward_indices_pointer, out=forward_indices_pointer) np.cumsum(reverse_indices_pointer, out=reverse_indices_pointer) forward_indices = np.memmap( paths.forward_indices, dtype=np.int32, mode="w+", shape=(forward_indices_pointer[-1],) ) reverse_indices = np.memmap( paths.reverse_indices, dtype=np.int32, mode="w+", shape=(reverse_indices_pointer[-1],) ) forward_cursor = forward_indices_pointer.copy() reverse_cursor = reverse_indices_pointer.copy() for u_str, v_str in tqdm( edges(), total=n_edges, unit="edge", unit_scale=True, desc="filling edges", disable=not progress, ): u = node_to_id[u_str] v = node_to_id[v_str] forward_indices[forward_cursor[u]] = v forward_cursor[u] += 1 reverse_indices[reverse_cursor[v]] = u reverse_cursor[v] += 1 forward_indices_pointer.tofile(paths.forward_indices_pointer) reverse_indices_pointer.tofile(paths.reverse_indices_pointer) reverse_indices.flush() forward_indices.flush() return GraphCache(paths) cthoyt-pystow-54a586e/src/pystow/impl.py000066400000000000000000002256311525577235500204320ustar00rootroot00000000000000"""Module implementation.""" from __future__ import annotations import bz2 import gzip import io import json import logging import lzma import os import pickle import sqlite3 import tarfile from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import closing, contextmanager from io import BytesIO, StringIO from pathlib import Path from typing import IO, TYPE_CHECKING, Any, Literal, TypeAlias, cast, overload from . import utils from .constants import JSON, Provider from .utils import ( base_from_gzip_name, download_from_google, download_from_s3, get_base, get_mode_pair, gunzip, mkdir, name_from_s3_key, name_from_url, open_tarfile, open_zipfile, path_to_sqlite, read_tarfile_csv, read_tarfile_xml, read_zip_np, read_zipfile_csv, ) from .utils.download import DownloadKwargs from .utils.safe_open import zstd_open if TYPE_CHECKING: from pathlib import PurePath import botocore.client import bs4 import lxml.etree import numpy import pandas as pd import rdflib __all__ = [ "Module", "VersionHint", ] logger = logging.getLogger(__name__) #: A type hint for something that can be passed to the #: `version` argument of Module.join, Module.ensure, etc. VersionHint: TypeAlias = str | Callable[[], str | None] | None class Module: """The class wrapping the directory lookup implementation.""" def __init__(self, base: str | Path, ensure_exists: bool = True) -> None: """Initialize the module. :param base: The base directory for the module :param ensure_exists: Should the base directory be created automatically? Defaults to true. """ self.base = Path(base) mkdir(self.base, ensure_exists=ensure_exists) @classmethod def from_key(cls, key: str, *subkeys: str, ensure_exists: bool = True) -> Module: """Get a module for the given directory or one of its subdirectories. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param ensure_exists: Should all directories be created automatically? Defaults to true. :returns: A module """ base = get_base(key, ensure_exists=False) rv = cls(base=base, ensure_exists=ensure_exists) if subkeys: rv = rv.module(*subkeys, ensure_exists=ensure_exists) return rv def module(self, *subkeys: str, ensure_exists: bool = True) -> Module: """Get a module for a subdirectory of the current module. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param ensure_exists: Should all directories be created automatically? Defaults to true. :returns: A module representing the subdirectory based on the given ``subkeys``. """ base = self.join(*subkeys, ensure_exists=False) return Module(base=base, ensure_exists=ensure_exists) # docstr-coverage:excused `overload` @overload def join( self, *subkeys: str, name: str | None = ..., ensure_exists: bool = ..., version: VersionHint = ..., return_version: Literal[True], ) -> tuple[Path, str | None]: ... # docstr-coverage:excused `overload` @overload def join( self, *subkeys: str, name: str | None = ..., ensure_exists: bool = ..., version: VersionHint = ..., return_version: Literal[False], ) -> Path: ... # docstr-coverage:excused `overload` @overload def join( self, *subkeys: str, name: str | None = ..., ensure_exists: bool = ..., version: VersionHint = ..., ) -> Path: ... # When return_version is omitted (defaults to False) def join( self, *subkeys: str, name: str | None = None, ensure_exists: bool = True, version: VersionHint = None, return_version: bool = False, ) -> Path | tuple[Path, str | None]: """Get a subdirectory of the current module. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param ensure_exists: Should all directories be created automatically? Defaults to true. :param name: The name of the file (optional) inside the folder :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. The following example describes how to store the versioned data from the Rhea database for biologically relevant chemical reactions. .. code-block:: import pystow import requests def get_rhea_version() -> str: res = requests.get("https://ftp.expasy.org/databases/rhea/rhea-release.properties") _, _, version = res.text.splitlines()[0].partition("=") return version # Assume you want to download the data from # ftp://ftp.expasy.org/databases/rhea/rdf/rhea.rdf.gz, make a path # with the same name module = pystow.module("rhea") path = module.join(name="rhea.rdf.gz", version=get_rhea_version) :param return_version: If true, returns the processed version :returns: The path of the directory or subdirectory for the given module. """ rv: Path = self.base if mod_version := self._resolve_version(version): subkeys = (mod_version, *subkeys) if subkeys: rv = rv.joinpath(*subkeys) mkdir(rv, ensure_exists=ensure_exists) if name: rv = rv.joinpath(name) if return_version: return rv, mod_version return rv @classmethod def _resolve_version(cls, version: VersionHint) -> str | None: match version: case None: return None case str(): cls._raise_for_invalid_version(version) return version case _: # if the version is given as a no-argument callable, # then it should be called and a version is returned rv: str | None = version() if rv is None: return None cls._raise_for_invalid_version(rv) return rv @staticmethod def _raise_for_invalid_version(version: str) -> None: if "/" in version or os.sep in version: raise ValueError( f"slashes and `{os.sep}` not allowed in versions because of " f"conflicts with file path construction: {version}" ) def joinpath_sqlite(self, *subkeys: str, name: str) -> str: """Get an SQLite database connection string. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the database file. :returns: A SQLite path string. """ path = self.join(*subkeys, name=name, ensure_exists=True) return path_to_sqlite(path) def ensure( self, *subkeys: str, url: str, name: str | None = None, version: VersionHint = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, ) -> Path: """Ensure a file is downloaded. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. The following example describes how to store the versioned data from the Rhea database for biologically relevant chemical reactions. .. code-block:: import pystow import requests def get_rhea_version() -> str: res = requests.get("https://ftp.expasy.org/databases/rhea/rhea-release.properties") _, _, version = res.text.splitlines()[0].partition("=") return version module = pystow.module("rhea") path = module.ensure( url="ftp://ftp.expasy.org/databases/rhea/rdf/rhea.rdf.gz", version=get_rhea_version, ) :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :returns: The path of the file that has been downloaded (or already exists) """ if name is None: name = name_from_url(url) path, version = self.join( *subkeys, name=name, version=version, ensure_exists=True, return_version=True ) utils.download( url=url, path=path, force=force, _version=version, **(download_kwargs or {}), ) return path def ensure_soup( self, *subkeys: str, url: str, name: str | None = None, version: VersionHint = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rt", "w", "wt", "rb", "wb"] = "r", open_kwargs: Mapping[str, Any] | None = None, beautiful_soup_kwargs: Mapping[str, Any] | None = None, ) -> bs4.BeautifulSoup: """Ensure a webpage is downloaded and parsed with :mod:`BeautifulSoup`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param beautiful_soup_kwargs: Additional keyword arguments passed to :class:`BeautifulSoup` :returns: An BeautifulSoup object .. note:: If you don't need to cache, consider using :func:`pystow.utils.get_soup` instead. """ from bs4 import BeautifulSoup beautiful_soup_kwargs = dict(beautiful_soup_kwargs or {}) beautiful_soup_kwargs.setdefault("features", "html.parser") with self.ensure_open( *subkeys, url=url, name=name, version=version, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, ) as file: soup = BeautifulSoup(file, **beautiful_soup_kwargs) return soup def ensure_custom( self, *subkeys: str, name: str, force: bool = False, provider: Provider, **kwargs: Any, ) -> Path: """Ensure a file is present, and run a custom create function otherwise. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The file name. :param force: Should the file be re-created, even if the path already exists? :param provider: The file provider. Will be run with the path as the first positional argument, if the file needs to be generated. :param kwargs: Additional keyword-based parameters passed to the provider. :returns: The path of the file that has been created (or already exists) :raises ValueError: If the provider was called but the file was not created by it. """ path = self.join(*subkeys, name=name, ensure_exists=True) if path.is_file() and not force: return path provider(path, **kwargs) if not path.is_file(): raise ValueError(f"Provider {provider} did not create the file at {path}!") return path def ensure_untar( self, *subkeys: str, url: str, name: str | None = None, directory: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, extract_kwargs: Mapping[str, Any] | None = None, ) -> Path: """Ensure a tar file is downloaded and unarchived. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param directory: Overrides the name of the directory into which the tar archive is extracted. If none given, will use the stem of the file name that gets downloaded. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param extract_kwargs: Keyword arguments to pass to :meth:`tarfile.TarFile.extract_all`. :returns: The path of the directory where the file that has been downloaded gets extracted to """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) if directory is None: # rhea-rxn.tar.gz -> rhea-rxn suffixes_len = sum(len(suffix) for suffix in path.suffixes) directory = path.name[:-suffixes_len] unzipped_path = path.parent.joinpath(directory) if unzipped_path.is_dir() and not force: return unzipped_path unzipped_path.mkdir(exist_ok=True, parents=True) with tarfile.open(path) as tar_file: tar_file.extractall(unzipped_path, **(extract_kwargs or {})) # noqa:S202 return unzipped_path def ensure_gunzip( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, autoclean: bool = True, download_kwargs: DownloadKwargs | None = None, ) -> Path: """Ensure a tar.gz file is downloaded and unarchived. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param autoclean: Should the zipped file be deleted? :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :returns: The path of the directory where the file that has been downloaded gets extracted to """ if name is None: name = name_from_url(url) gunzipped_name = base_from_gzip_name(name) gunzipped_path = self.join(*subkeys, name=gunzipped_name, ensure_exists=True) if gunzipped_path.is_file() and not force: return gunzipped_path path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, ) gunzip(path, gunzipped_path) if autoclean: logger.info("removing original gzipped file %s", path) path.unlink() return gunzipped_path # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open( self, *subkeys: str, url: str, name: str | None, version: VersionHint = ..., force: bool, download_kwargs: DownloadKwargs | None, mode: Literal["r", "rt", "w", "wt"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[StringIO]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open( self, *subkeys: str, url: str, name: str | None, version: VersionHint = ..., force: bool, download_kwargs: DownloadKwargs | None, mode: Literal["rb", "wb"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[BytesIO]: ... @contextmanager def ensure_open( self, *subkeys: str, url: str, name: str | None = None, version: VersionHint = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rt", "w", "wt", "rb", "wb"] = "r", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[StringIO | BytesIO]: """Ensure a file is downloaded and open it. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :yields: An open file object """ path = self.ensure( *subkeys, url=url, name=name, version=version, force=force, download_kwargs=download_kwargs, ) open_kwargs = {} if open_kwargs is None else dict(open_kwargs) open_kwargs.setdefault("mode", mode) with path.open(**open_kwargs) as file: yield file @staticmethod def _raise_for_mode_ensure_mismatch( mode: Literal["r", "rt", "w", "wt", "wb", "rb"], ensure_exists: bool ) -> None: """Raise an exception for a mismatch between the mode and the ensure status. :param mode: The file opening mode :param ensure_exists: The value for ensuring a file :raises ValueError: In the following situations: 1. If the file should be opened in write mode, and it is not ensured to exist 2. If the file should be opened in read mode, and it is ensured to exist. This is bad because it will create a file when there previously wasn't one """ if "w" in mode and not ensure_exists: raise ValueError if "r" in mode and ensure_exists: raise ValueError # docstr-coverage:excused `overload` @overload @contextmanager def open( self, *subkeys: str, name: str, mode: Literal["r", "rt", "w", "wt"] = ..., open_kwargs: Mapping[str, Any] | None = None, ensure_exists: bool, ) -> Generator[StringIO]: ... # docstr-coverage:excused `overload` @overload @contextmanager def open( self, *subkeys: str, name: str, mode: Literal["rb", "wb"] = ..., open_kwargs: Mapping[str, Any] | None = None, ensure_exists: bool, ) -> Generator[BytesIO]: ... @contextmanager def open( self, *subkeys: str, name: str, mode: Literal["r", "rt", "w", "wt", "rb", "wb"] = "r", open_kwargs: Mapping[str, Any] | None = None, ensure_exists: bool = False, ) -> Generator[StringIO | BytesIO]: """Open a file. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param ensure_exists: Should the directory the file is in be made? Set to true on write operations. :yields: An open file object. This function should be called inside a context manager like in the following .. code-block:: python import pystow with pystow.module("test").open(name="test.tsv", mode="w") as file: print("Test text!", file=file) """ self._raise_for_mode_ensure_mismatch(mode, ensure_exists) path = self.join(*subkeys, name=name, ensure_exists=ensure_exists) open_kwargs = {} if open_kwargs is None else dict(open_kwargs) open_kwargs.setdefault("mode", mode) with path.open(**open_kwargs) as file: yield file # docstr-coverage:excused `overload` @overload @contextmanager def open_gz( self, *subkeys: str, name: str, mode: Literal["r", "w", "rt", "wt"] = ..., open_kwargs: Mapping[str, Any] | None, ensure_exists: bool, ) -> Generator[StringIO]: ... # docstr-coverage:excused `overload` @overload @contextmanager def open_gz( self, *subkeys: str, name: str, mode: Literal["rb", "wb"] = ..., open_kwargs: Mapping[str, Any] | None, ensure_exists: bool, ) -> Generator[BytesIO]: ... @contextmanager def open_gz( self, *subkeys: str, name: str, mode: Literal["r", "w", "rt", "wt", "rb", "wb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, ensure_exists: bool = False, ) -> Generator[StringIO | BytesIO]: """Open a gzipped file that exists already. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param mode: The read mode, passed to :func:`gzip.open` :param open_kwargs: Additional keyword arguments passed to :func:`gzip.open` :param ensure_exists: Should the file be made? Set to true on write operations. :yields: An open file object """ self._raise_for_mode_ensure_mismatch(mode, ensure_exists) path = self.join(*subkeys, name=name, ensure_exists=ensure_exists) open_kwargs = {} if open_kwargs is None else dict(open_kwargs) open_kwargs.setdefault("mode", mode) with gzip.open(path, **open_kwargs) as file: yield file # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_lzma( self, *subkeys: str, url: str, name: str | None, force: bool, download_kwargs: DownloadKwargs | None, mode: Literal["r", "w", "rt", "wt"] = "rt", open_kwargs: Mapping[str, Any] | None, ) -> Generator[io.TextIOWrapper[lzma.LZMAFile]]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_lzma( self, *subkeys: str, url: str, name: str | None, force: bool, download_kwargs: DownloadKwargs | None, mode: Literal["rb", "wb"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[lzma.LZMAFile]: ... @contextmanager def ensure_open_lzma( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rb", "w", "wb", "rt", "wt"] = "rt", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[lzma.LZMAFile | io.TextIOWrapper[lzma.LZMAFile]]: """Ensure a LZMA-compressed file is downloaded and open a file inside it. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`lzma.open` :param open_kwargs: Additional keyword arguments passed to :func:`lzma.open` :yields: An open file object """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) open_kwargs = {} if open_kwargs is None else dict(open_kwargs) open_kwargs.setdefault("mode", mode) with lzma.open(path, **open_kwargs) as file: yield file @contextmanager def ensure_open_zstd( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rb", "w", "wb", "rt", "wt"] = "rt", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[io.BufferedIOBase]: """Ensure a zSTD-compressed file is downloaded and open a file inside it. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`zstd.open` :param open_kwargs: Additional keyword arguments passed to :func:`zstd.open` :yields: An open file object """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) open_kwargs = {} if open_kwargs is None else dict(open_kwargs) open_kwargs.setdefault("mode", mode) with zstd_open(path, **open_kwargs) as file: yield file # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_tarfile( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = ..., force: bool = ..., download_kwargs: DownloadKwargs | None = ..., mode: Literal["rt"] = ..., open_kwargs: Mapping[str, Any] | None = ..., ) -> Generator[IO[str]]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_tarfile( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = ..., force: bool = ..., download_kwargs: DownloadKwargs | None = ..., mode: Literal["r", "rb"] = ..., open_kwargs: Mapping[str, Any] | None = ..., ) -> Generator[IO[bytes]]: ... @contextmanager def ensure_open_tarfile( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rb", "rt"] = "r", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[IO[str]] | Generator[IO[bytes]]: """Ensure a tar file is downloaded and open a file inside it. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`tarfile.open` :param open_kwargs: Additional keyword arguments passed to :func:`tarfile.open` :yields: An open file object """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) open_kwargs = {} if open_kwargs is None else dict(open_kwargs) open_kwargs.setdefault("mode", mode) operation, representation = get_mode_pair(open_kwargs.pop("mode"), interpretation="binary") with open_tarfile( path=path, inner_path=inner_path, operation=operation, representation=representation, open_kwargs=open_kwargs, ) as file: yield file # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_zip( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = ..., force: bool = ..., download_kwargs: DownloadKwargs | None = ..., mode: Literal["r", "rb"] = ..., open_kwargs: Mapping[str, Any] | None = ..., ) -> Generator[IO[bytes]]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_zip( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = ..., force: bool = ..., download_kwargs: DownloadKwargs | None = ..., mode: Literal["rt"] = ..., open_kwargs: Mapping[str, Any] | None = ..., ) -> Generator[IO[str]]: ... @contextmanager def ensure_open_zip( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rb", "rt"] = "r", zipfile_kwargs: Mapping[str, Any] | None = None, open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[IO[str]] | Generator[IO[bytes]]: """Ensure a file is downloaded then open it with :mod:`zipfile`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`zipfile.open`. Defaults to bytes mode for ``r`` and ``w``. :param zipfile_kwargs: Additional keyword arguments passed to :class:`zipfile.ZipFile` :param open_kwargs: Additional keyword arguments passed to :func:`zipfile.open` :yields: An open file object """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) open_kwargs = {} if open_kwargs is None else dict(open_kwargs) open_kwargs.setdefault("mode", mode) operation, representation = get_mode_pair(open_kwargs.pop("mode"), interpretation="binary") with open_zipfile( path=path, inner_path=inner_path, operation=operation, representation=representation, zipfile_kwargs=zipfile_kwargs, open_kwargs=open_kwargs, ) as file: yield file # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_gz( self, *subkeys: str, url: str, name: str | None, force: bool, download_kwargs: DownloadKwargs | None, mode: Literal["r", "w", "rt", "wt"] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[StringIO]: ... # docstr-coverage:excused `overload` @overload @contextmanager def ensure_open_gz( self, *subkeys: str, url: str, name: str | None, force: bool, download_kwargs: DownloadKwargs | None, mode: Literal[ "rb", "wb", ] = ..., open_kwargs: Mapping[str, Any] | None, ) -> Generator[BytesIO]: ... @contextmanager def ensure_open_gz( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, version: VersionHint = None, download_kwargs: DownloadKwargs | None = None, mode: Literal["r", "rb", "w", "wb", "rt", "wt"] = "rb", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[StringIO | BytesIO]: """Ensure a gzipped file is downloaded and open a file inside it. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`gzip.open` :param open_kwargs: Additional keyword arguments passed to :func:`gzip.open` :yields: An open file object """ path = self.ensure( *subkeys, url=url, name=name, force=force, version=version, download_kwargs=download_kwargs, ) open_kwargs = {} if open_kwargs is None else dict(open_kwargs) open_kwargs.setdefault("mode", mode) with gzip.open(path, **open_kwargs) as file: yield file @contextmanager def ensure_open_bz2( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, version: VersionHint = None, download_kwargs: DownloadKwargs | None = None, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, ) -> Generator[bz2.BZ2File]: """Ensure a BZ2-compressed file is downloaded and open a file inside it. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`bz2.open` :param open_kwargs: Additional keyword arguments passed to :func:`bz2.open` :yields: An open file object """ path = self.ensure( *subkeys, url=url, name=name, force=force, version=version, download_kwargs=download_kwargs, ) open_kwargs = {} if open_kwargs is None else dict(open_kwargs) open_kwargs.setdefault("mode", mode) with bz2.open(path, **open_kwargs) as file: yield file def ensure_csv( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, version: VersionHint = None, download_kwargs: DownloadKwargs | None = None, read_csv_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Download a CSV and open as a dataframe with :mod:`pandas`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param read_csv_kwargs: Keyword arguments to pass through to :func:`pandas.read_csv`. :returns: A pandas DataFrame """ import pandas as pd path = self.ensure( *subkeys, url=url, name=name, force=force, version=version, download_kwargs=download_kwargs, ) return pd.read_csv(path, **_clean_csv_kwargs(read_csv_kwargs)) def load_df( self, *subkeys: str, name: str, read_csv_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Open a pre-existing CSV as a dataframe with :mod:`pandas`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param read_csv_kwargs: Keyword arguments to pass through to :func:`pandas.read_csv`. :returns: A pandas DataFrame """ import pandas as pd with self.open(*subkeys, name=name, mode="r", ensure_exists=False) as file: return pd.read_csv(file, **_clean_csv_kwargs(read_csv_kwargs)) def dump_df( self, *subkeys: str, name: str, obj: pd.DataFrame, sep: str = "\t", index: bool = False, to_csv_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump a dataframe to a TSV file with :mod:`pandas`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param obj: The dataframe to dump :param sep: The separator to use, defaults to a tab :param index: Should the index be dumped? Defaults to false. :param to_csv_kwargs: Keyword arguments to pass through to :meth:`pandas.DataFrame.to_csv`. """ to_csv_kwargs = {} if to_csv_kwargs is None else dict(to_csv_kwargs) to_csv_kwargs.setdefault("sep", sep) to_csv_kwargs.setdefault("index", index) # should this use unified opener instead? Pandas is pretty smart... path = self.join(*subkeys, name=name) obj.to_csv(path, **to_csv_kwargs) def ensure_json( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, version: VersionHint = None, download_kwargs: DownloadKwargs | None = None, open_kwargs: Mapping[str, Any] | None = None, json_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Download JSON and open with :mod:`json`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param open_kwargs: Additional keyword arguments passed to :func:`open` :param json_load_kwargs: Keyword arguments to pass through to :func:`json.load`. :returns: A JSON object (list, dict, etc.) """ with self.ensure_open( *subkeys, url=url, name=name, force=force, version=version, download_kwargs=download_kwargs, open_kwargs=open_kwargs, ) as file: return json.load(file, **(json_load_kwargs or {})) def ensure_json_bz2( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, version: VersionHint = None, download_kwargs: DownloadKwargs | None = None, open_kwargs: Mapping[str, Any] | None = None, json_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Download BZ2-compressed JSON and open with :mod:`json`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param version: The optional version, or no-argument callable that returns an optional version. This is prepended before the subkeys. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param open_kwargs: Additional keyword arguments passed to :func:`bz2.open` :param json_load_kwargs: Keyword arguments to pass through to :func:`json.load`. :returns: A JSON object (list, dict, etc.) """ with self.ensure_open_bz2( *subkeys, url=url, name=name, force=force, version=version, download_kwargs=download_kwargs, open_kwargs=open_kwargs, ) as file: return json.load(file, **(json_load_kwargs or {})) def ensure_yaml( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, open_kwargs: Mapping[str, Any] | None = None, yaml_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Download YAML and open with :mod:`yaml`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param open_kwargs: Additional keyword arguments passed to :func:`open` :param yaml_load_kwargs: Keyword arguments to pass through to :func:`yaml.safe_load`. :returns: A JSON object (list, dict, etc.) """ import yaml with self.ensure_open( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, open_kwargs=open_kwargs, ) as file: return yaml.safe_load(file, **(yaml_load_kwargs or {})) def load_yaml( self, *subkeys: str, name: str, open_kwargs: Mapping[str, Any] | None = None, yaml_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Open a JSON file :mod:`json`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param open_kwargs: Additional keyword arguments passed to :func:`open` :param yaml_load_kwargs: Keyword arguments to pass through to :func:`yaml.safe_load`. :returns: A JSON object (list, dict, etc.) """ import yaml with self.open( *subkeys, name=name, mode="r", open_kwargs=open_kwargs, ensure_exists=False ) as file: return yaml.safe_load(file, **(yaml_load_kwargs or {})) def load_json( self, *subkeys: str, name: str, open_kwargs: Mapping[str, Any] | None = None, json_load_kwargs: Mapping[str, Any] | None = None, ) -> JSON: """Open a JSON file :mod:`json`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param open_kwargs: Additional keyword arguments passed to :func:`open` :param json_load_kwargs: Keyword arguments to pass through to :func:`json.load`. :returns: A JSON object (list, dict, etc.) """ with self.open( *subkeys, name=name, mode="r", open_kwargs=open_kwargs, ensure_exists=False ) as file: return json.load(file, **(json_load_kwargs or {})) def dump_json( self, *subkeys: str, name: str, obj: JSON, open_kwargs: Mapping[str, Any] | None = None, json_dump_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump an object to a file with :mod:`json`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param obj: The object to dump :param open_kwargs: Additional keyword arguments passed to :func:`open` :param json_dump_kwargs: Keyword arguments to pass through to :func:`json.dump`. """ with self.open( *subkeys, name=name, mode="w", open_kwargs=open_kwargs, ensure_exists=True ) as file: json.dump(obj, file, **(json_dump_kwargs or {})) def ensure_pickle( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, pickle_load_kwargs: Mapping[str, Any] | None = None, ) -> Any: """Download a pickle file and open with :mod:`pickle`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param pickle_load_kwargs: Keyword arguments to pass through to :func:`pickle.load`. :returns: Any object """ with self.ensure_open( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, ) as file: return pickle.load(file, **(pickle_load_kwargs or {})) def load_pickle( self, *subkeys: str, name: str, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, pickle_load_kwargs: Mapping[str, Any] | None = None, ) -> Any: """Open a pickle file with :mod:`pickle`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param pickle_load_kwargs: Keyword arguments to pass through to :func:`pickle.load`. :returns: Any object """ with self.open( *subkeys, name=name, mode=mode, open_kwargs=open_kwargs, ensure_exists=False, ) as file: return pickle.load(file, **(pickle_load_kwargs or {})) def dump_pickle( self, *subkeys: str, name: str, obj: Any, mode: Literal["wb"] = "wb", open_kwargs: Mapping[str, Any] | None = None, pickle_dump_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump an object to a file with :mod:`pickle`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param obj: The object to dump :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`open` :param pickle_dump_kwargs: Keyword arguments to pass through to :func:`pickle.dump`. """ with self.open( *subkeys, name=name, mode=mode, open_kwargs=open_kwargs, ensure_exists=True, ) as file: pickle.dump(obj, file, **(pickle_dump_kwargs or {})) def ensure_pickle_gz( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, pickle_load_kwargs: Mapping[str, Any] | None = None, ) -> Any: """Download a gzipped pickle file and open with :mod:`pickle`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param mode: The read mode, passed to :func:`gzip.open` :param open_kwargs: Additional keyword arguments passed to :func:`gzip.open` :param pickle_load_kwargs: Keyword arguments to pass through to :func:`pickle.load`. :returns: Any object """ with self.ensure_open_gz( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs, mode=mode, open_kwargs=open_kwargs, ) as file: return pickle.load(file, **(pickle_load_kwargs or {})) def load_pickle_gz( self, *subkeys: str, name: str, mode: Literal["rb"] = "rb", open_kwargs: Mapping[str, Any] | None = None, pickle_load_kwargs: Mapping[str, Any] | None = None, ) -> Any: """Open a gzipped pickle file with :mod:`pickle`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param mode: The read mode, passed to :func:`open` :param open_kwargs: Additional keyword arguments passed to :func:`gzip.open` :param pickle_load_kwargs: Keyword arguments to pass through to :func:`pickle.load`. :returns: Any object """ with self.open_gz( *subkeys, name=name, mode=mode, open_kwargs=open_kwargs, ensure_exists=False, ) as file: return pickle.load(file, **(pickle_load_kwargs or {})) def ensure_excel( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, read_excel_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Download an excel file and open as a dataframe with :mod:`pandas`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param read_excel_kwargs: Keyword arguments to pass through to :func:`pandas.read_excel`. :returns: A pandas DataFrame """ import pandas as pd path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) return pd.read_excel(path, **(read_excel_kwargs or {})) def ensure_tar_df( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, read_csv_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Download a tar file and open an inner file as a dataframe with :mod:`pandas`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param read_csv_kwargs: Keyword arguments to pass through to :func:`pandas.read_csv`. :returns: A dataframe .. warning:: If you have lots of files to read in the same archive, it's better just to unzip first. """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) return read_tarfile_csv( path=path, inner_path=inner_path, **_clean_csv_kwargs(read_csv_kwargs) ) def ensure_xml( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, parse_kwargs: Mapping[str, Any] | None = None, ) -> lxml.etree.ElementTree: """Download an XML file and open it with :mod:`lxml`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param parse_kwargs: Keyword arguments to pass through to :func:`lxml.etree.parse`. :returns: An ElementTree object .. warning:: If you have lots of files to read in the same archive, it's better just to unzip first. """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) return utils.read_xml(path, **(parse_kwargs or {})) def load_xml( self, *subkeys: str, name: str, parse_kwargs: Mapping[str, Any] | None = None, ) -> lxml.etree.ElementTree: """Load an XML file with :mod:`lxml`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param parse_kwargs: Keyword arguments to pass through to :func:`lxml.etree.parse`. :returns: An ElementTree object .. warning:: If you have lots of files to read in the same archive, it's better just to unzip first. """ from lxml import etree with self.open(*subkeys, mode="r", name=name, ensure_exists=False) as file: return etree.parse(file, **(parse_kwargs or {})) def dump_xml( self, *subkeys: str, name: str, obj: lxml.etree.ElementTree, open_kwargs: Mapping[str, Any] | None = None, write_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump an XML element tree to a file with :mod:`lxml`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param obj: The object to dump :param open_kwargs: Additional keyword arguments passed to :func:`open` :param write_kwargs: Keyword arguments to pass through to :func:`lxml.etree.ElementTree.write`. """ with self.open( *subkeys, name=name, mode="wb", open_kwargs=open_kwargs, ensure_exists=True ) as file: obj.write(file, **(write_kwargs or {})) def ensure_tar_xml( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, parse_kwargs: Mapping[str, Any] | None = None, ) -> lxml.etree.ElementTree: """Download a tar file and open an inner file as an XML with :mod:`lxml`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param parse_kwargs: Keyword arguments to pass through to :func:`lxml.etree.parse`. :returns: An ElementTree object .. warning:: If you have lots of files to read in the same archive, it's better just to unzip first. """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) return read_tarfile_xml(path=path, inner_path=inner_path, **(parse_kwargs or {})) def ensure_zip_df( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, read_csv_kwargs: Mapping[str, Any] | None = None, ) -> pd.DataFrame: """Download a zip file and open an inner file as a dataframe with :mod:`pandas`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param read_csv_kwargs: Keyword arguments to pass through to :func:`pandas.read_csv`. :returns: A pandas DataFrame """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) return read_zipfile_csv( path=path, inner_path=inner_path, **_clean_csv_kwargs(read_csv_kwargs) ) def ensure_zip_np( self, *subkeys: str, url: str, inner_path: str | PurePath, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, load_kwargs: Mapping[str, Any] | None = None, ) -> numpy.typing.ArrayLike: """Download a zip file and open an inner file as an array-like with :mod:`numpy`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param inner_path: The relative path to the file inside the archive :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param load_kwargs: Additional keyword arguments that are passed through to :func:`read_zip_np` and transitively to :func:`numpy.load`. :returns: An array-like object """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) return read_zip_np(path=path, inner_path=inner_path, **(load_kwargs or {})) def ensure_rdf( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, precache: bool = True, parse_kwargs: Mapping[str, Any] | None = None, ) -> rdflib.Graph: """Download a RDF file and open with :mod:`rdflib`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :param precache: Should the parsed :class:`rdflib.Graph` be stored as a pickle for fast loading? :param parse_kwargs: Keyword arguments to pass through to :func:`pystow.utils.read_rdflib` and transitively to :func:`rdflib.Graph.parse`. :returns: An RDF graph """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) if not precache: return utils.read_rdflib(path=path, **(parse_kwargs or {})) cache_path = path.with_suffix(path.suffix + ".pickle.gz") if cache_path.exists() and not force: import rdflib with gzip.open(cache_path, "rb") as file: return cast(rdflib.Graph, pickle.load(file)) rv = utils.read_rdflib(path=path, **(parse_kwargs or {})) with gzip.open(cache_path, "wb") as file: pickle.dump(rv, file, protocol=pickle.HIGHEST_PROTOCOL) return rv def load_rdf( self, *subkeys: str, name: str | None = None, parse_kwargs: Mapping[str, Any] | None = None, ) -> rdflib.Graph: """Open an RDF file with :mod:`rdflib`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param parse_kwargs: Keyword arguments to pass through to :func:`pystow.utils.read_rdflib` and transitively to :func:`rdflib.Graph.parse`. :returns: An RDF graph """ path = self.join(*subkeys, name=name, ensure_exists=False) return utils.read_rdflib(path=path, **(parse_kwargs or {})) def dump_rdf( self, *subkeys: str, name: str, obj: rdflib.Graph, format: str = "turtle", serialize_kwargs: Mapping[str, Any] | None = None, ) -> None: """Dump an RDF graph to a file with :mod:`rdflib`. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file to open :param obj: The object to dump :param format: The format to dump in :param serialize_kwargs: Keyword arguments to through to :func:`rdflib.Graph.serialize`. """ path = self.join(*subkeys, name=name, ensure_exists=False) serialize_kwargs = {} if serialize_kwargs is None else dict(serialize_kwargs) serialize_kwargs.setdefault("format", format) obj.serialize(path, **serialize_kwargs) def ensure_from_s3( self, *subkeys: str, s3_bucket: str, s3_key: str | Sequence[str], name: str | None = None, client: botocore.client.BaseClient | None = None, client_kwargs: Mapping[str, Any] | None = None, download_file_kwargs: Mapping[str, Any] | None = None, force: bool = False, ) -> Path: """Ensure a file is downloaded. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param s3_bucket: The S3 bucket name :param s3_key: The S3 key name :param name: Overrides the name of the file at the end of the S3 key, if given. :param client: A botocore client. If none given, one will be created automatically :param client_kwargs: Keyword arguments to be passed to the client on instantiation. :param download_file_kwargs: Keyword arguments to be passed to :func:`boto3.s3.transfer.S3Transfer.download_file` :param force: Should the download be done again, even if the path already exists? Defaults to false. :returns: The path of the file that has been downloaded (or already exists) """ if not isinstance(s3_key, str): s3_key = "/".join(s3_key) # join sequence if name is None: name = name_from_s3_key(s3_key) path = self.join(*subkeys, name=name, ensure_exists=True) download_from_s3( s3_bucket=s3_bucket, s3_key=s3_key, path=path, client=client, client_kwargs=client_kwargs, force=force, download_file_kwargs=download_file_kwargs, ) return path def ensure_from_google( self, *subkeys: str, name: str, file_id: str, force: bool = False, download_kwargs: Mapping[str, Any] | None = None, ) -> Path: """Ensure a file is downloaded from Google Drive. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param name: The name of the file :param file_id: The file identifier of the Google file. If your share link is https://drive.google.com/file/d/1AsPPU4ka1Rc9u-XYMGWtvV65hF3egi0z/view, then your file ID is ``1AsPPU4ka1Rc9u-XYMGWtvV65hF3egi0z``. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download_from_google`. :returns: The path of the file that has been downloaded (or already exists) """ path = self.join(*subkeys, name=name, ensure_exists=True) download_from_google(file_id, path, force=force, **(download_kwargs or {})) return path @contextmanager def ensure_open_sqlite( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, ) -> Generator[sqlite3.Connection]: """Ensure and connect to a SQLite database. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :yields: An instance of :class:`sqlite3.Connection` from :func:`sqlite3.connect` Example usage: >>> import pystow >>> import pandas as pd >>> url = "https://s3.amazonaws.com/bbop-sqlite/obi.db" >>> sql = "SELECT * FROM entailed_edge LIMIT 10" >>> module = pystow.module("test") >>> with module.ensure_open_sqlite(url=url) as conn: >>> df = pd.read_sql(sql, conn) """ path = self.ensure( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) with closing(sqlite3.connect(path.as_posix())) as conn: yield conn @contextmanager def ensure_open_sqlite_gz( self, *subkeys: str, url: str, name: str | None = None, force: bool = False, download_kwargs: DownloadKwargs | None = None, ) -> Generator[sqlite3.Connection]: """Ensure and connect to a SQLite database that's gzipped. Unfortunately, it's a paid feature to directly read gzipped sqlite files, so this automatically gunzips it first. :param subkeys: A sequence of additional strings to join. If none are given, returns the directory for this module. :param url: The URL to download. :param name: Overrides the name of the file at the end of the URL, if given. Also useful for URLs that don't have proper filenames with extensions. :param force: Should the download be done again, even if the path already exists? Defaults to false. :param download_kwargs: Keyword arguments to pass through to :func:`pystow.utils.download`. :yields: An instance of :class:`sqlite3.Connection` from :func:`sqlite3.connect` Example usage: >>> import pystow >>> import pandas as pd >>> url = "https://s3.amazonaws.com/bbop-sqlite/hp.db.gz" >>> module = pystow.module("test") >>> sql = "SELECT * FROM entailed_edge LIMIT 10" >>> with module.ensure_open_sqlite_gz(url=url) as conn: >>> df = pd.read_sql(sql, conn) """ path = self.ensure_gunzip( *subkeys, url=url, name=name, force=force, download_kwargs=download_kwargs ) with closing(sqlite3.connect(path.as_posix())) as conn: yield conn def _clean_csv_kwargs(read_csv_kwargs: Mapping[str, Any] | None) -> dict[str, Any]: read_csv_kwargs = {} if read_csv_kwargs is None else dict(read_csv_kwargs) read_csv_kwargs.setdefault("sep", "\t") return read_csv_kwargs cthoyt-pystow-54a586e/src/pystow/py.typed000066400000000000000000000000001525577235500205730ustar00rootroot00000000000000cthoyt-pystow-54a586e/src/pystow/utils/000077500000000000000000000000001525577235500202465ustar00rootroot00000000000000cthoyt-pystow-54a586e/src/pystow/utils/__init__.py000066400000000000000000001417231525577235500223670ustar00rootroot00000000000000"""Utilities.""" from __future__ import annotations import contextlib import csv import gzip import io import logging import lzma import pickle import shutil import tarfile import typing import warnings import zipfile from collections.abc import Callable, Generator, Iterable, Mapping, Sequence from io import BytesIO, StringIO from pathlib import Path, PurePath, PurePosixPath from subprocess import check_output from typing import ( IO, TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, TypeVar, cast, overload, ) from urllib.parse import urlparse from uuid import uuid4 from tqdm.auto import tqdm from .download import ( DownloadBackend, DownloadError, DownloadKwargs, UnexpectedDirectoryError, download, download_from_google, download_from_s3, ) from .env import ( get_base, get_home, get_name, getenv_path, mkdir, mock_envvar, mock_home, use_appdirs, ) from .hashing import ( Hash, HexDigestError, HexDigestMismatch, get_hash_hexdigest, get_hashes, get_hexdigests_remote, get_offending_hexdigests, raise_on_digest_mismatch, ) from .io_typing import ( _OPERATION_TO_UNQUALIFIED_MODE, MODE_MAP, OPERATION_VALUES, REPRESENTATION_VALUES, REVERSE_MODE_MAP, InvalidOperationError, InvalidRepresentationError, Operation, Reader, Representation, Writer, ensure_sensible_default_encoding, ensure_sensible_newline, get_mode_pair, ) from .iter import reyield from .pydantic_utils import ( ModelValidateFailureAction, iter_pydantic_jsonl, iter_pydantic_tsv, model_dump_yaml, read_pydantic_json, read_pydantic_jsonl, read_pydantic_tsv, read_pydantic_yaml, stream_write_pydantic_jsonl, write_pydantic_json, write_pydantic_jsonl, write_pydantic_yaml, ) from .safe_open import ( _wrap_binary_if_needed, is_url, open_inner_zipfile, open_url, safe_open, safe_open_dict_reader, safe_open_json, safe_open_yaml, safe_read_text, safe_write_text, write_json, write_yaml, ) from .testing import requires_package from ..constants import README_TEXT, TimeoutHint if TYPE_CHECKING: import bs4 import lxml.etree import numpy.typing import pandas import rdflib __all__ = [ "MODE_MAP", "OPERATION_VALUES", "REPRESENTATION_VALUES", "REVERSE_MODE_MAP", "ArchiveType", "DownloadBackend", "DownloadError", "DownloadKwargs", "Hash", "HeaderMismatchError", "HexDigestError", "HexDigestMismatch", "InvalidOperationError", "InvalidRepresentationError", "ModelValidateFailureAction", "Operation", "Reader", "Representation", "UnexpectedDirectory", "UnexpectedDirectoryError", "Writer", "download", "download_from_google", "download_from_s3", "get_base", "get_commit", "get_df_io", "get_hash_hexdigest", "get_hashes", "get_hexdigests_remote", "get_home", "get_mode_pair", "get_name", "get_np_io", "get_offending_hexdigests", "get_soup", "get_xml_root", "getenv_path", "gunzip", "gzip_compress", "is_url", "iter_pydantic_jsonl", "iter_pydantic_tsv", "iter_tarred_csvs", "iter_tarred_files", "iter_zipped_csvs", "iter_zipped_files", "iterparse_xml", "mkdir", "mock_envvar", "mock_home", "model_dump_yaml", "n", "name_from_s3_key", "name_from_url", "open_archive", "open_inner_tarfile", "open_inner_zipfile", "open_tarfile", "open_url", "open_zip_dict_reader", "open_zip_reader", "open_zip_writer", "open_zipfile", "path_to_sqlite", "raise_on_digest_mismatch", "read_lzma_csv", "read_pydantic_json", "read_pydantic_jsonl", "read_pydantic_tsv", "read_pydantic_yaml", "read_rdf", "read_rdflib", "read_tarfile_csv", "read_tarfile_xml", "read_xml", "read_zip_np", "read_zipfile_csv", "read_zipfile_rdf", "read_zipfile_xml", "requires_package", "reyield", "safe_open", "safe_open_dict_reader", "safe_open_dict_writer", "safe_open_json", "safe_open_reader", "safe_open_writer", "safe_open_yaml", "safe_read_text", "safe_tarfile_open", "safe_write_text", "safe_zipfile_open", "stream_write_pydantic_jsonl", "tarfile_writestr", "use_appdirs", "write_json", "write_lzma_csv", "write_pickle_gz", "write_pydantic_json", "write_pydantic_jsonl", "write_pydantic_yaml", "write_rdflib", "write_tarfile_csv", "write_tarfile_xml", "write_yaml", "write_zipfile_csv", "write_zipfile_np", "write_zipfile_rdf", "write_zipfile_xml", ] logger = logging.getLogger(__name__) #: Backwards compatible name UnexpectedDirectory = UnexpectedDirectoryError def name_from_url(url: str) -> str: """Get the filename from the end of the URL. :param url: A URL :returns: The name of the file at the end of the URL """ parse_result = urlparse(url) path = PurePosixPath(parse_result.path) name = path.name return name def base_from_gzip_name(name: str) -> str: """Get the base name for a file after stripping the gz ending. :param name: The name of the gz file :returns: The cleaned name of the file, with no gz ending :raises ValueError: if the file does not end with ".gz" """ if not name.endswith(".gz"): raise ValueError(f"Name does not end with .gz: {name}") return name[: -len(".gz")] def name_from_s3_key(key: str) -> str: """Get the filename from the S3 key. :param key: A S3 path :returns: The name of the file """ return key.split("/")[-1] def n() -> str: """Get a random string for testing. :returns: A random string for testing purposes. """ return str(uuid4()) def get_df_io(df: pandas.DataFrame, sep: str = "\t", index: bool = False, **kwargs: Any) -> BytesIO: """Get the dataframe as bytes. :param df: A dataframe :param sep: The separator in the dataframe. Overrides Pandas default to use a tab. :param index: Should the index be output? Overrides the Pandas default to be false. :param kwargs: Additional kwargs to pass to :func:`pandas.DataFrame.to_csv`. :returns: A bytes object that can be used as a file. """ return io.BytesIO(df.to_csv(sep=sep, index=index, **kwargs).encode("utf-8")) def get_np_io(arr: numpy.typing.ArrayLike, **kwargs: Any) -> BytesIO: """Get the numpy object as bytes. :param arr: Array-like :param kwargs: Additional kwargs to pass to :func:`numpy.save`. :returns: A bytes object that can be used as a file. """ import numpy as np bio = BytesIO() np.save(bio, arr, **kwargs) bio.seek(0) return bio def write_pickle_gz( obj: Any, path: str | Path, **kwargs: Any, ) -> None: """Write an object to a gzipped pickle. :param obj: The object to write :param path: The path of the file to write to :param kwargs: Additional kwargs to pass to :func:`pickle.dump` """ with safe_open(path, representation="binary", operation="write") as file: pickle.dump(obj, file, **kwargs) def write_lzma_csv( df: pandas.DataFrame, path: str | Path, sep: str = "\t", index: bool = False, **kwargs: Any, ) -> None: """Write a dataframe as an lzma-compressed file. :param df: A dataframe :param path: The path to the resulting LZMA compressed dataframe file :param sep: The separator in the dataframe. Overrides Pandas default to use a tab. :param index: Should the index be output? Overrides the Pandas default to be false. :param kwargs: Additional kwargs to pass to :func:`get_df_io` and transitively to :func:`pandas.DataFrame.to_csv`. """ with lzma.open(path, "wb") as file: df.to_csv(file, sep=sep, index=index, **kwargs) def read_lzma_csv( path: str | Path, sep: str = "\t", **kwargs: Any, ) -> pandas.DataFrame: """Read a dataframe from a lzma-compressed file. :param path: The path to the resulting LZMA compressed dataframe file :param sep: The separator in the dataframe. Overrides Pandas default to use a tab. :param kwargs: Additional kwargs to pass to :func:`get_df_io` and transitively to :func:`pandas.DataFrame.to_csv`. """ import pandas as pd with lzma.open(path, "rb") as file: return pd.read_csv(file, sep=sep, **kwargs) def write_zipfile_csv( df: pandas.DataFrame, path: str | Path, inner_path: str | PurePath, sep: str = "\t", index: bool = False, **kwargs: Any, ) -> None: """Write a dataframe to an inner CSV file to a zip archive. :param df: A dataframe :param path: The path to the resulting zip archive :param inner_path: The path inside the zip archive to write the dataframe :param sep: The separator in the dataframe. Overrides Pandas default to use a tab. :param index: Should the index be output? Overrides the Pandas default to be false. :param kwargs: Additional kwargs to pass to :func:`get_df_io` and transitively to :func:`pandas.DataFrame.to_csv`. """ with open_zipfile(path, inner_path, operation="write", representation="binary") as file: df.to_csv(file, sep=sep, index=index, **kwargs) def read_zipfile_csv( path: str | Path, inner_path: str | PurePath, sep: str = "\t", **kwargs: Any ) -> pandas.DataFrame: """Read an inner CSV file from a zip archive. :param path: The path to the zip archive :param inner_path: The path inside the zip archive to the dataframe :param sep: The separator in the dataframe. Overrides Pandas default to use a tab. :param kwargs: Additional kwargs to pass to :func:`pandas.read_csv`. :returns: A dataframe """ import pandas as pd with open_zipfile(path, inner_path, representation="text", operation="read") as file: return pd.read_csv(file, sep=sep, **kwargs) # docstr-coverage:excused `overload` @typing.overload @contextlib.contextmanager def open_zipfile( path: str | Path, inner_path: str | PurePath, *, operation: Operation = ..., representation: Literal["text"] = ..., zipfile_kwargs: Mapping[str, Any] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[str]]: ... # docstr-coverage:excused `overload` @typing.overload @contextlib.contextmanager def open_zipfile( path: str | Path, inner_path: str | PurePath, *, operation: Operation = ..., representation: Literal["binary"] = ..., zipfile_kwargs: Mapping[str, Any] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[bytes]]: ... @contextlib.contextmanager def open_zipfile( path: str | Path, inner_path: str | PurePath, *, operation: Operation = "read", representation: Representation = "text", zipfile_kwargs: Mapping[str, Any] | None = None, open_kwargs: Mapping[str, Any] | None = None, encoding: str | None = None, newline: str | None = None, ) -> Generator[IO[str]] | Generator[IO[bytes]]: """Open a zipfile.""" mode = _OPERATION_TO_UNQUALIFIED_MODE[operation] with ( zipfile.ZipFile(file=path, mode=mode, **(zipfile_kwargs or {})) as zip_file, open_inner_zipfile( zip_file, inner_path, operation=operation, representation=representation, open_kwargs=open_kwargs, encoding=encoding, newline=newline, ) as file, ): yield file # docstr-coverage:excused `overload` @overload @contextlib.contextmanager def open_tarfile( path: str | Path, inner_path: str | PurePath, *, operation: Operation = ..., representation: Literal["binary"] = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[bytes]]: ... # docstr-coverage:excused `overload` @overload @contextlib.contextmanager def open_tarfile( path: str | Path, inner_path: str | PurePath, *, operation: Operation = ..., representation: Literal["text"] = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[str]]: ... @contextlib.contextmanager def open_tarfile( path: str | Path, inner_path: str | PurePath, *, operation: Operation = "read", representation: Representation = "binary", open_kwargs: Mapping[str, Any] | None = None, encoding: str | None = None, newline: str | None = None, ) -> Generator[IO[bytes]] | Generator[IO[str]]: """Open a tar file.""" mode = _OPERATION_TO_UNQUALIFIED_MODE[operation] with ( tarfile.open(path, mode, **(open_kwargs or {})) as tar_file, open_inner_tarfile( tar_file, inner_path, operation=operation, representation=representation, encoding=encoding, newline=newline, ) as file, ): yield file # docstr-coverage:excused `overload` @overload @contextlib.contextmanager def open_inner_tarfile( tar_file: tarfile.TarFile, inner_path: str | PurePath, *, operation: Operation = ..., representation: Literal["text"] = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[str]]: ... # docstr-coverage:excused `overload` @overload @contextlib.contextmanager def open_inner_tarfile( tar_file: tarfile.TarFile, inner_path: str | PurePath, *, operation: Operation = ..., representation: Literal["binary"] = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[bytes]]: ... @contextlib.contextmanager def open_inner_tarfile( tar_file: tarfile.TarFile, inner_path: str | PurePath, *, operation: Operation = "read", representation: Representation = "text", encoding: str | None = None, newline: str | None = None, ) -> Generator[IO[str]] | Generator[IO[bytes]]: """Open an inner tar file.""" # by convention, paths inside tar files are POSIX-style, even on Windows inner_path = str(PurePosixPath(inner_path)) encoding = ensure_sensible_default_encoding(encoding, representation=representation) newline = ensure_sensible_newline(newline, representation=representation) if operation == "read": member = tar_file.getmember(inner_path) file = tar_file.extractfile(member) if file is None: raise FileNotFoundError(f"could not find {inner_path} in tarfile {tar_file}") with _wrap_binary_if_needed(file, representation, encoding=encoding, newline=newline) as yf: yield yf elif operation == "write": if representation == "binary": file = BytesIO() yield file file.seek(0) elif representation == "text": sio = StringIO(newline=newline) yield sio sio.seek(0) file = BytesIO(sio.getvalue().encode(cast(str, encoding))) else: raise InvalidRepresentationError(representation) tarinfo = tarfile.TarInfo(name=inner_path) tarinfo.size = len(file.getbuffer()) tar_file.addfile(tarinfo, file) else: raise InvalidOperationError(operation) @contextlib.contextmanager def open_zip_reader( path: str | Path, inner_path: str | PurePath, delimiter: str = "\t", **kwargs: Any ) -> Generator[Reader]: """Read an inner CSV file from a zip archive. :param path: The path to the zip archive :param inner_path: The path inside the zip archive to the CSV :param delimiter: The separator in the CSV. Defaults to tab. :param kwargs: Additional kwargs to pass to :func:`csv.reader`. :returns: A reader over the file """ with open_zipfile(path, inner_path, representation="text") as file: yield csv.reader(file, delimiter=delimiter, **kwargs) @contextlib.contextmanager def open_zip_dict_reader( path: str | Path, inner_path: str | PurePath, delimiter: str = "\t", **kwargs: Any ) -> Generator[csv.DictReader[str]]: """Read an inner CSV file from a zip archive. :param path: The path to the zip archive :param inner_path: The path inside the zip archive to the CSV :param delimiter: The separator in the CSV. Defaults to tab. :param kwargs: Additional kwargs to pass to :class:`csv.DictReader`. :returns: A dictionary reader over the file """ with open_zipfile(path, inner_path, representation="text") as file: yield csv.DictReader(file, delimiter=delimiter, **kwargs) @contextlib.contextmanager def open_zip_writer( path: str | Path, inner_path: str | PurePath, delimiter: str = "\t", **kwargs: Any ) -> Generator[Writer]: """Open a writer for an inner CSV file from a zip archive. :param path: The path to the zip archive :param inner_path: The path inside the zip archive to the CSV :param delimiter: The separator in the CSV. Defaults to tab. :param kwargs: Additional kwargs to pass to :func:`csv.writer`. :returns: A writer over the file """ with open_zipfile(path, inner_path, operation="write", representation="text") as file: yield csv.writer(file, delimiter=delimiter, **kwargs) def write_zipfile_xml( element_tree: lxml.etree.ElementTree, path: str | Path, inner_path: str | PurePath, **kwargs: Any, ) -> None: """Write an XML element tree to an inner XML file to a zip archive. :param element_tree: An XML element tree :param path: The path to the resulting zip archive :param inner_path: The path inside the zip archive to write the XML element :param kwargs: Additional kwargs to pass to :func:`lxml.etree.tostring` """ from lxml import etree kwargs.setdefault("pretty_print", True) with open_zipfile(path, inner_path, operation="write", representation="binary") as file: file.write(etree.tostring(element_tree, **kwargs)) def read_xml( path: str | Path, open_kwargs: dict[str, Any] | None = None, **kwargs: Any ) -> lxml.etree.ElementTree: """Read an XML element tree. :param path: The path to an XML file :param kwargs: Additional kwargs to pass to :func:`lxml.etree.parse`. :returns: An element tree """ from lxml import etree with safe_open(path, representation="binary") as file: return etree.parse(file, **kwargs) def get_xml_root(path: str | Path, **kwargs: Any) -> lxml.etree.Element: """Read an XML element tree then get the root element. :param path: The path to an XML file :param kwargs: Additional kwargs to pass to :func:`lxml.etree.parse`. :returns: The root element from the element tree """ return read_xml(path, **kwargs).getroot() def iterparse_xml(path: str | Path, tag: str | tuple[str, ...], **kwargs: Any) -> Iterable[Any]: """Parse the given tags, iteratively.""" from lxml import etree with safe_open(path, representation="binary") as file: yield from etree.iterparse(file, tag=tag, **kwargs) def read_zipfile_xml( path: str | Path, inner_path: str | PurePath, **kwargs: Any ) -> lxml.etree.ElementTree: """Read an inner XML file from a zip archive. :param path: The path to the zip archive :param inner_path: The path inside the zip archive to the xml file :param kwargs: Additional kwargs to pass to :func:`lxml.etree.parse` :returns: An XML element tree """ from lxml import etree with open_zipfile(path, inner_path, operation="read", representation="binary") as file: return etree.parse(file, **kwargs) def write_zipfile_np( arr: numpy.typing.ArrayLike, path: str | Path, inner_path: str | PurePath, **kwargs: Any, ) -> None: """Write a dataframe to an inner CSV file to a zip archive. :param arr: Array-like :param path: The path to the resulting zip archive :param inner_path: The path inside the zip archive to write the dataframe :param kwargs: Additional kwargs to pass to :func:`get_np_io` and transitively to :func:`numpy.save`. """ import numpy as np with open_zipfile(path, inner_path, operation="write", representation="binary") as file: np.save(file, arr, **kwargs) def read_zip_np( path: str | Path, inner_path: str | PurePath, **kwargs: Any ) -> numpy.typing.ArrayLike: """Read an inner numpy array-like from a zip archive. :param path: The path to the zip archive :param inner_path: The path inside the zip archive to the dataframe :param kwargs: Additional kwargs to pass to :func:`numpy.load`. :returns: A numpy array or other object """ import numpy as np with open_zipfile(path, inner_path, operation="read", representation="binary") as file: return cast(np.typing.ArrayLike, np.load(file, **kwargs)) def read_zipfile_rdf(path: str | Path, inner_path: str | PurePath, **kwargs: Any) -> rdflib.Graph: """Read an inner RDF file from a zip archive. :param path: The path to the zip archive :param inner_path: The path inside the zip archive to the dataframe :param kwargs: Additional kwargs to pass to :meth:`rdflib.Graph.parse`. :returns: A graph """ import rdflib graph = rdflib.Graph() with open_zipfile(path, inner_path, operation="read", representation="binary") as file: graph.parse(file, **kwargs) return graph def write_zipfile_rdf( graph: rdflib.Graph, path: str | Path, inner_path: str | PurePath, **kwargs: Any ) -> None: """Read an inner RDF file from a zip archive. :param graph: The graph to write :param path: The path to the zip archive :param inner_path: The path inside the zip archive to the dataframe :param kwargs: Additional kwargs to pass to :meth:`rdflib.Graph.parse`. """ with open_zipfile(path, inner_path, operation="write", representation="binary") as file: graph.serialize(file, **kwargs) def write_tarfile_csv( df: pandas.DataFrame, path: str | Path, inner_path: str | PurePath, sep: str = "\t", index: bool = False, **kwargs: Any, ) -> None: """Write a dataframe to an inner CSV file from a tar archive. :param df: A dataframe :param path: The path to the resulting tar archive :param inner_path: The path inside the tar archive to write the dataframe :param sep: The separator in the dataframe. Overrides Pandas default to use a tab. :param index: Should the index be output? Overrides the Pandas default to be false. :param kwargs: Additional kwargs to pass to :func:`get_df_io` and transitively to :func:`pandas.DataFrame.to_csv`. """ with open_tarfile(path, inner_path, operation="write") as file: df.to_csv(file, sep=sep, index=index, **kwargs) def write_tarfile_xml( element_tree: lxml.etree.ElementTree, path: str | Path, inner_path: str | PurePath, **kwargs: Any, ) -> None: """Write an XML document a tar archive. :param element_tree: An element :param path: The path to the resulting tar archive :param inner_path: The path inside the tar archive to write the dataframe :param kwargs: Additional kwargs to pass to :func:`lxml.etree.tostring` """ from lxml import etree kwargs.setdefault("pretty_print", True) with open_tarfile(path, inner_path, operation="write") as file: file.write(etree.tostring(element_tree, **kwargs)) def read_tarfile_csv( path: str | Path, inner_path: str | PurePath, sep: str = "\t", **kwargs: Any ) -> pandas.DataFrame: """Read an inner CSV file from a tar archive. :param path: The path to the tar archive :param inner_path: The path inside the tar archive to the dataframe :param sep: The separator in the dataframe. Overrides Pandas default to use a tab. :param kwargs: Additional kwargs to pass to :func:`pandas.read_csv`. :returns: A dataframe """ import pandas as pd with open_tarfile(path, inner_path) as file: return pd.read_csv(file, sep=sep, **kwargs) def read_tarfile_xml( path: str | Path, inner_path: str | PurePath, **kwargs: Any ) -> lxml.etree.ElementTree: """Read an inner XML file from a tar archive. :param path: The path to the tar archive :param inner_path: The path inside the tar archive to the xml file :param kwargs: Additional kwargs to pass to :func:`lxml.etree.parse` :returns: An XML element tree """ from lxml import etree with open_tarfile(path, inner_path) as file: return etree.parse(file, **kwargs) def read_rdf(path: str | Path, **kwargs: Any) -> rdflib.Graph: """Read an RDF file with :mod:`rdflib` via :func:`read_rdflib`.""" warnings.warn( "use read_rdflib() instead - this new function has a more precise name", DeprecationWarning, stacklevel=2, ) return read_rdflib(path, **kwargs) def read_rdflib(path: str | Path, **kwargs: Any) -> rdflib.Graph: """Read an RDF file with :mod:`rdflib`. :param path: The path to the RDF file :param kwargs: Additional kwargs to pass to :func:`rdflib.Graph.parse` :returns: A parsed RDF graph """ import rdflib graph = rdflib.Graph() with safe_open(path, representation="binary", operation="read") as file: graph.parse(file, **kwargs) return graph def write_rdflib( graph: rdflib.Graph, path: str | Path | IO[str] | IO[bytes], *, format: str | None = None, ) -> None: """Write an RDF file with :mod:`rdflib`. :param graph: The RDF graph :param path: The path to the RDF file :param format: The format to write the RDF to. Defaults to ttl """ with safe_open(path, representation="binary", operation="write") as file: graph.serialize(file, format=format or "ttl") def write_sql(df: pandas.DataFrame, name: str, path: str | Path, **kwargs: Any) -> None: """Write a dataframe as a SQL table. :param df: A dataframe :param name: The table the database to write to :param path: The path to the resulting tar archive :param kwargs: Additional keyword arguments to pass to :meth:`pandas.DataFrame.to_sql` """ import sqlite3 with contextlib.closing(sqlite3.connect(path)) as conn: df.to_sql(name, conn, **kwargs) def get_commit(org: str, repo: str, provider: str = "git") -> str: """Get last commit hash for the given repo. :param org: The GitHub organization or owner :param repo: The GitHub repository name :param provider: The method for getting the most recent commit :returns: A commit hash's hex digest as a string :raises ValueError: if an invalid provider is given """ if provider == "git": output = check_output(["git", "ls-remote", f"https://github.com/{org}/{repo}"]) # noqa lines = (line.strip().split("\t") for line in output.decode("utf8").splitlines()) rv = next(line[0] for line in lines if line[1] == "HEAD") elif provider == "github": url = f"https://api.github.com/repos/{org}/{repo}/branches/master" res_json = safe_open_json(url) rv = res_json["commit"]["sha"] else: raise ValueError(f"invalid implementation: {provider}") return rv def ensure_readme() -> None: """Ensure there's a README in the PyStow data directory. :raises PermissionError: If the script calling this function does not have adequate permissions to write a file into the PyStow home directory. """ try: readme_path = get_home(ensure_exists=True).joinpath("README.md") except PermissionError as e: raise PermissionError( "PyStow was not able to create its home directory in due to a lack of " "permissions. This can happen, e.g., if you're working on a server where you don't " "have full rights. See https://pystow.readthedocs.io/en/latest/installation.html#" "configuration for instructions on choosing a different home folder location for " "PyStow to somewhere where you have write permissions." ) from e if readme_path.is_file(): return with readme_path.open("w", encoding="utf8") as file: print(README_TEXT, file=file) def path_to_sqlite(path: str | Path) -> str: """Convert a path to a SQLite connection string. :param path: A path to a SQLite database file :returns: A standard connection string to the database """ path = Path(path).expanduser().resolve() return f"sqlite:///{path.as_posix()}" def gunzip(source: str | Path, target: str | Path | None = None, *, cleanup: bool = False) -> Path: """Unzip a file in the source to the target. :param source: The path to an input file :param target: The path to an output file :param cleanup: Whether to clean the output file """ source = Path(source).expanduser().resolve() if target is None: raise NotImplementedError else: target = Path(target).expanduser().resolve() with gzip.open(source, "rb") as in_file, open(target, "wb") as out_file: shutil.copyfileobj(in_file, out_file) if cleanup: source.unlink() return target def gzip_compress( source: str | Path, *, target: str | Path | None = None, cleanup: bool = False ) -> Path: """Compress a file, then delete the original.""" source = Path(source).expanduser().resolve() if target is None: target = source.with_suffix(source.suffix + ".gz") else: target = Path(target).expanduser().resolve() with open(source, "rb") as in_file, gzip.open(target, "wb") as out_file: shutil.copyfileobj(in_file, out_file) if cleanup: source.unlink() return target class BatchedWriter: """Wrap a writer for batching.""" writer: Writer batch_size: int batch: list[Iterable[Any]] def __init__(self, writer: Writer, batch_size: int) -> None: """Initialize the batched writer.""" self.writer = writer self.batch_size = batch_size self.batch = [] def writerow(self, row: Iterable[Any], /) -> None: """Write a single row to the batch.""" self.batch.append(row) if len(self.batch) >= self.batch_size: self.flush() def writerows(self, rows: Iterable[Iterable[Any]], /) -> None: """Write multiple rows to the batch.""" # TODO get fancy and only fill up the batch # then batch on the rows themselves self.batch.extend(rows) if len(self.batch) >= self.batch_size: self.flush() def flush(self) -> None: """Write the rest of the remaining batch.""" self.writer.writerows(self.batch) self.batch.clear() @contextlib.contextmanager def safe_open_writer( f: str | Path | IO[str], *, delimiter: str = "\t", buffering: int | None = None, batch_size: int | None = None, **kwargs: Any, ) -> Generator[Writer]: """Open a CSV writer, wrapping :func:`csv.writer`. :param f: A path to a file, or an already open text-based IO object :param delimiter: The delimiter for writing to CSV :param buffering: The buffer size for the file. If not given, defaults to -1, which opens in unbuffered mode :param batch_size: The number of rows to write in each batch. If you're using this, it's probably also good to add ``buffering=1024*1024`` to the file opener :param kwargs: Keyword arguments to pass to :func:`csv.writer` :yields: A CSV writer object, constructed from :func:`csv.writer` """ with safe_open(f, operation="write", representation="text", buffering=buffering) as file: writer = csv.writer(file, delimiter=delimiter, **kwargs) if batch_size is not None: batched_writer = BatchedWriter(writer, batch_size) # cast because the batched writer object is a duck yield cast(Writer, batched_writer) # make sure there's a final flush of any remaining rows batched_writer.flush() else: yield writer @contextlib.contextmanager def safe_open_dict_writer( f: str | Path | IO[str], fieldnames: typing.Sequence[str], *, delimiter: str = "\t", **kwargs: Any, ) -> Generator[csv.DictWriter[str]]: """Open a CSV dictionary writer, wrapping :func:`csv.DictWriter`. :param f: A path to a file, or an already open text-based IO object :param fieldnames: A path to a file, or an already open text-based IO object :param delimiter: The delimiter for writing to CSV :param kwargs: Keyword arguments to pass to :func:`csv.DictWriter` :yields: A CSV dictionary writer object, constructed from :func:`csv.DictWriter` """ with safe_open(f, operation="write", representation="text", newline="") as file: yield csv.DictWriter(file, fieldnames, delimiter=delimiter, **kwargs) @contextlib.contextmanager def safe_open_reader( f: str | Path | IO[str], *, delimiter: str = "\t", **kwargs: Any ) -> Generator[Reader]: """Open a CSV reader, wrapping :func:`csv.reader`. :param f: A path to a file, or an already open text-based IO object :param delimiter: The delimiter for writing to CSV :param kwargs: Keyword arguments to pass to :func:`csv.reader` :yields: A CSV reader object, constructed from :func:`csv.reader` """ with safe_open(f, operation="read", representation="text", newline="") as file: yield csv.reader(file, delimiter=delimiter, **kwargs) def get_soup( url: str, *, verify: bool = True, timeout: TimeoutHint | None = None, user_agent: str | None = None, ) -> bs4.BeautifulSoup: """Get a beautiful soup parsed version of the given web page. :param url: The URL to download and parse with BeautifulSoup :param verify: Should SSL be used? This is almost always true, except for Ensembl, which makes a big pain :param timeout: How many integer seconds to wait for a response? Defaults to 15 if none given. :param user_agent: A custom user-agent to set, e.g., to avoid anti-crawling mechanisms :returns: A BeautifulSoup object """ import requests from bs4 import BeautifulSoup headers = {} if user_agent: headers["User-Agent"] = user_agent res = requests.get(url, verify=verify, timeout=timeout or 15, headers=headers) soup = BeautifulSoup(res.text, features="html.parser") return soup ArchiveType_contra = TypeVar("ArchiveType_contra", contravariant=True) ArchiveInfo_co = TypeVar("ArchiveInfo_co", covariant=True) Predicate: TypeAlias = Callable[[ArchiveInfo_co], bool] class ArchivedFileIterator(Protocol[ArchiveType_contra, ArchiveInfo_co]): """A protocol for opening files in an archive.""" # docstr-coverage:excused `overload` @overload def __call__( self, path: str | Path | ArchiveType_contra, *, representation: Literal["binary"] = ..., progress: bool = ..., tqdm_kwargs: Mapping[str, Any] | None = ..., keep: Predicate[ArchiveInfo_co] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Iterable[IO[bytes]]: ... # docstr-coverage:excused `overload` @overload def __call__( self, path: str | Path | ArchiveType_contra, *, representation: Literal["text"] = ..., progress: bool = ..., tqdm_kwargs: Mapping[str, Any] | None = ..., keep: Predicate[ArchiveInfo_co] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Iterable[IO[str]]: ... def __call__( self, path: str | Path | ArchiveType_contra, *, representation: Representation = ..., progress: bool = True, tqdm_kwargs: Mapping[str, Any] | None = ..., keep: Predicate[ArchiveInfo_co] | None = ..., open_kwargs: Mapping[str, Any] | None = None, encoding: str | None = ..., newline: str | None = ..., ) -> Iterable[IO[str]] | Iterable[IO[bytes]]: ... # docstr-coverage:excused `overload` @overload def iter_tarred_files( path: str | Path | tarfile.TarFile, *, representation: Literal["binary"] = ..., progress: bool = ..., tqdm_kwargs: Mapping[str, Any] | None = ..., keep: Predicate[tarfile.TarInfo] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Iterable[IO[bytes]]: ... # docstr-coverage:excused `overload` @overload def iter_tarred_files( path: str | Path | tarfile.TarFile, *, representation: Literal["text"] = ..., progress: bool = ..., tqdm_kwargs: Mapping[str, Any] | None = ..., keep: Predicate[tarfile.TarInfo] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Iterable[IO[str]]: ... def iter_tarred_files( path: str | Path | tarfile.TarFile, *, representation: Representation = "text", progress: bool = True, tqdm_kwargs: Mapping[str, Any] | None = None, keep: Predicate[tarfile.TarInfo] | None = None, open_kwargs: Mapping[str, Any] | None = None, encoding: str | None = None, newline: str | None = None, ) -> Iterable[IO[str]] | Iterable[IO[bytes]]: """Iterate over opened files in a tar archive in read mode.""" encoding = ensure_sensible_default_encoding(encoding, representation=representation) newline = ensure_sensible_newline(newline, representation=representation) with safe_tarfile_open(path) as tar_file: _tqdm_kwargs: dict[str, Any] = { "unit": "file", "unit_scale": True, } if isinstance(tar_file.name, str | Path): _tqdm_kwargs["desc"] = f"reading {Path(tar_file.name).name}" if tqdm_kwargs is not None: _tqdm_kwargs.update(tqdm_kwargs) for member in tqdm(tar_file.getmembers(), disable=not progress, **_tqdm_kwargs): if keep is not None and not keep(member): continue match tar_file.extractfile(member, **(open_kwargs or {})): case None: continue case io.BufferedReader() as reader: if representation == "text": yield io.TextIOWrapper(reader, encoding=encoding, newline=newline) else: yield reader @contextlib.contextmanager def safe_tarfile_open( tar_file: str | Path | tarfile.TarFile, ) -> Generator[tarfile.TarFile]: """Open a tar archive safely.""" if isinstance(tar_file, str | Path): with tarfile.open(Path(tar_file).expanduser().resolve(), mode="r") as yv: yield yv else: yield tar_file ReturnType: TypeAlias = Literal["sequence", "record"] # docstr-coverage:excused `overload` @overload def iter_tarred_csvs( path: str | Path | tarfile.TarFile, *, progress: bool = ..., return_type: Literal["sequence"] = ..., max_line_length: int | None = ..., ) -> Iterable[Sequence[str]]: ... # docstr-coverage:excused `overload` @overload def iter_tarred_csvs( path: str | Path | tarfile.TarFile, *, progress: bool = ..., return_type: Literal["record"] = ..., max_line_length: int | None = ..., ) -> Iterable[dict[str, Any]]: ... def iter_tarred_csvs( path: str | Path | tarfile.TarFile, *, progress: bool = True, return_type: ReturnType = "sequence", tqdm_kwargs: Mapping[str, Any] | None = None, max_line_length: int | None = None, encoding: str | None = None, ) -> Iterable[Sequence[str]] | Iterable[dict[str, Any]]: """Iterate over the lines from tarred CSV files.""" yield from _iter_archived_csvs( path, progress=progress, return_type=return_type, iter_files=iter_tarred_files, keep=_keep_tar_info_csv, tqdm_kwargs=tqdm_kwargs, max_line_length=max_line_length, encoding=encoding, ) def _keep_tar_info_csv(tar_info: tarfile.TarInfo) -> bool: return tar_info.name.endswith(".csv") # docstr-coverage:excused `overload` @overload def iter_zipped_files( path: str | Path | zipfile.ZipFile, *, representation: Literal["binary"] = ..., progress: bool = ..., tqdm_kwargs: Mapping[str, Any] | None = ..., keep: Predicate[zipfile.ZipInfo] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Iterable[IO[bytes]]: ... # docstr-coverage:excused `overload` @overload def iter_zipped_files( path: str | Path | zipfile.ZipFile, *, representation: Literal["text"] = ..., progress: bool = ..., tqdm_kwargs: Mapping[str, Any] | None = ..., keep: Predicate[zipfile.ZipInfo] | None = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Iterable[IO[str]]: ... def iter_zipped_files( path: str | Path | zipfile.ZipFile, *, representation: Representation = "text", progress: bool = True, tqdm_kwargs: Mapping[str, Any] | None = None, keep: Predicate[zipfile.ZipInfo] | None = None, open_kwargs: Mapping[str, Any] | None = None, encoding: str | None = None, newline: str | None = None, ) -> Iterable[IO[str]] | Iterable[IO[bytes]]: """Iterate over opened files in a zip file in read mode.""" with safe_zipfile_open(path) as zip_file: _tqdm_kwargs = { "desc": f"reading {zip_file.filename}", "unit": "file", "unit_scale": True, } if tqdm_kwargs is not None: _tqdm_kwargs.update(tqdm_kwargs) for info in tqdm(zip_file.infolist(), disable=not progress, **_tqdm_kwargs): if keep is not None and not keep(info): continue with open_inner_zipfile( zip_file, info.filename, operation="read", representation=representation, open_kwargs=open_kwargs, encoding=encoding, newline=newline, ) as file: yield file @contextlib.contextmanager def safe_zipfile_open( zip_file: str | Path | zipfile.ZipFile, ) -> Generator[zipfile.ZipFile]: """Open a zip archive safely.""" if isinstance(zip_file, str | Path): with zipfile.ZipFile(Path(zip_file).expanduser().resolve(), mode="r") as yv: yield yv else: yield zip_file # docstr-coverage:excused `overload` @overload def iter_zipped_csvs( path: str | Path | zipfile.ZipFile, *, progress: bool = ..., return_type: Literal["sequence"] = ..., tqdm_kwargs: Mapping[str, Any] | None = ..., max_line_length: int | None = ..., ) -> Iterable[Sequence[str]]: ... # docstr-coverage:excused `overload` @overload def iter_zipped_csvs( path: str | Path | zipfile.ZipFile, *, progress: bool = ..., return_type: Literal["record"] = ..., tqdm_kwargs: Mapping[str, Any] | None = ..., max_line_length: int | None = ..., ) -> Iterable[dict[str, Any]]: ... def iter_zipped_csvs( path: str | Path | zipfile.ZipFile, *, progress: bool = True, return_type: ReturnType = "sequence", tqdm_kwargs: Mapping[str, Any] | None = None, max_line_length: int | None = None, encoding: str | None = None, ) -> Iterable[Sequence[str]] | Iterable[dict[str, Any]]: """Iterate over the lines from zipped CSV files.""" yield from _iter_archived_csvs( path, progress=progress, return_type=return_type, iter_files=iter_zipped_files, keep=_keep_zip_info_csv, tqdm_kwargs=tqdm_kwargs, max_line_length=max_line_length, encoding=encoding, ) def _keep_zip_info_csv(zip_info: zipfile.ZipInfo) -> bool: return zip_info.filename.endswith(".csv") def _iter_archived_csvs( path: str | Path | ArchiveType_contra, *, progress: bool = True, tqdm_kwargs: Mapping[str, Any] | None = None, keep: Predicate[ArchiveInfo_co] | None = None, return_type: ReturnType = "sequence", iter_files: ArchivedFileIterator[ArchiveType_contra, ArchiveInfo_co], max_line_length: int | None = None, encoding: str | None = None, ) -> Iterable[Sequence[str]] | Iterable[dict[str, Any]]: """Iterate over the lines from zipped CSV files.""" header: Sequence[str] | None = None for file in iter_files( path, representation="text", progress=progress, tqdm_kwargs=tqdm_kwargs, keep=keep, encoding=encoding, newline="", ): filename = file.name if max_line_length is not None: # this will break everything if there's an issue in the # header, but we aren't going to consider that case it = _cut_long_lines(file, max_line_length, filename) else: it = file reader: csv.DictReader[str] | Reader match return_type: case "sequence": reader = csv.reader(it) case "record": reader = csv.DictReader(it) case _: raise ValueError(f"unrecognized return type {return_type}") if header is None: header = _get_header(reader) elif (current_header := _get_header(reader)) != header: raise HeaderMismatchError(header, current_header) rv = tqdm( reader, disable=not progress, leave=False, desc=f"reading {filename}", unit="row", unit_scale=True, ) yield from rv def _cut_long_lines(it: Iterable[str], max_length: int, name: str) -> Iterable[str]: for i, line in enumerate(it): if len(line) > max_length: tqdm.write(f"[{name}:{i:,}] line of length {len(line):,} is too long: {line[:100]}") continue yield line def _get_header(reader: csv.DictReader[str] | Reader) -> Sequence[str]: if isinstance(reader, csv.DictReader): return cast(Sequence[str], reader.fieldnames) else: return next(reader) class HeaderMismatchError(ValueError): """Raised when the current header in an archive of CSVs is different than the original.""" def __init__(self, original: Sequence[str], current: Sequence[str]) -> None: """Instantiate the error.""" self.original = original self.current = current def __str__(self) -> str: return ( f"header mismatch. first header was {self.original} " f"and current header is {self.current}" ) def tarfile_writestr(tar_file: tarfile.TarFile, filename: str, data: str) -> None: """Write text to a tarfile.""" # TODO later, combine with other tarfile writing tarfile_write_bytes(tar_file, filename, data.encode("utf-8")) def tarfile_write_bytes(tar_file: tarfile.TarFile, filename: str, data: bytes) -> None: """Write bytes to a tarfile.""" tar_info = tarfile.TarInfo(name=filename) tar_info.size = len(data) tar_file.addfile(tar_info, io.BytesIO(data)) #: The archive type ArchiveType: TypeAlias = Literal["tar", "zip"] # docstr-coverage:excused `overload` @overload @contextlib.contextmanager def open_archive( path: str | Path, inner_path: str | PurePath, *, archive_type: ArchiveType, operation: Operation = ..., representation: Literal["binary"] = ..., outer_open_kwargs: Mapping[str, Any] | None = ..., inner_open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[bytes]]: ... # docstr-coverage:excused `overload` @overload @contextlib.contextmanager def open_archive( path: str | Path, inner_path: str | PurePath, *, archive_type: ArchiveType, operation: Operation = ..., representation: Literal["text"] = ..., outer_open_kwargs: Mapping[str, Any] | None = ..., inner_open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[str]]: ... @contextlib.contextmanager def open_archive( path: str | Path, inner_path: str | PurePath, *, archive_type: ArchiveType, operation: Operation = "read", representation: Representation = "text", outer_open_kwargs: Mapping[str, Any] | None = None, inner_open_kwargs: Mapping[str, Any] | None = None, encoding: str | None = None, newline: str | None = None, ) -> Generator[IO[str]] | Generator[IO[bytes]]: """Open an archived file.""" if archive_type == "tar": with open_tarfile( path, inner_path, operation=operation, representation=representation, open_kwargs=outer_open_kwargs, encoding=encoding, newline=newline, ) as file: yield file elif archive_type == "zip": with open_zipfile( path, inner_path, operation=operation, representation=representation, zipfile_kwargs=outer_open_kwargs, open_kwargs=inner_open_kwargs, encoding=encoding, newline=newline, ) as file: yield file else: raise ValueError(f"unrecognized archive type: {archive_type}") cthoyt-pystow-54a586e/src/pystow/utils/download.py000066400000000000000000000340451525577235500224350ustar00rootroot00000000000000"""Download utilities.""" from __future__ import annotations import logging import shutil import urllib.error from collections.abc import Mapping from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypeAlias, TypedDict from urllib.request import urlretrieve from tqdm import tqdm from typing_extensions import Unpack from .hashing import raise_on_digest_mismatch from ..constants import TimeoutHint from ..version import VERSION if TYPE_CHECKING: import botocore.client import requests import requests.exceptions __all__ = [ "DownloadBackend", "DownloadError", "DownloadKwargs", "RequestKwargs", "UnexpectedDirectoryError", "download", "download_from_google", "download_from_s3", ] logger = logging.getLogger(__name__) #: Represents an available backend for downloading DownloadBackend: TypeAlias = Literal["urllib", "requests"] class TqdmReportHook(tqdm): # type:ignore """A custom progress bar that can be used with urllib. Based on https://gist.github.com/leimao/37ff6e990b3226c2c9670a2cd1e4a6f5 """ def update_to( self, blocks: int = 1, block_size: int = 1, total_size: int | None = None, ) -> None: """Update the internal state based on a urllib report hook. :param blocks: Number of blocks transferred so far :param block_size: Size of each block (in tqdm units) :param total_size: Total size (in tqdm units). If [default: None] remains unchanged. """ if total_size is not None: self.total = total_size self.update(blocks * block_size - self.n) # will also set self.n = b * bsize class RequestKwargs(TypedDict): """Keyword arguments for :func:`requests.get`.""" auth: NotRequired[tuple[str, str]] timeout: NotRequired[TimeoutHint] allow_redirects: NotRequired[bool] proxies: NotRequired[dict[str, str]] verify: NotRequired[bool] stream: NotRequired[bool] cert: NotRequired[str | tuple[str, str]] params: NotRequired[dict[str, Any]] headers: NotRequired[dict[str, str | bytes | None] | None] class DownloadKwargs(RequestKwargs): """Keyword arguments for :func:`download`.""" # note: `force` is intentionally omitted here because # it is passed through from other signature components clean_on_failure: NotRequired[bool] backend: NotRequired[DownloadBackend | None] hexdigests: NotRequired[Mapping[str, str] | None] hexdigests_remote: NotRequired[Mapping[str, str] | None] hexdigests_strict: NotRequired[bool] progress_bar: NotRequired[bool] tqdm_kwargs: NotRequired[Mapping[str, Any] | None] DEFAULT_AGENT = f"pystow v{VERSION}" def download( # noqa:C901 url: str, path: str | Path, *, force: bool = True, clean_on_failure: bool = True, backend: DownloadBackend | None = None, hexdigests: Mapping[str, str] | None = None, hexdigests_remote: Mapping[str, str] | None = None, hexdigests_strict: bool = False, progress_bar: bool = True, tqdm_kwargs: Mapping[str, Any] | None = None, _version: str | None = None, **kwargs: Unpack[RequestKwargs], ) -> None: """Download a file from a given URL. :param url: URL to download :param path: Path to download the file to :param force: If false and the file already exists, will not re-download. :param clean_on_failure: If true, will delete the file on any exception raised during download :param backend: The downloader to use. Choose 'urllib' or 'requests' :param hexdigests: The expected hexdigests as (algorithm_name, expected_hex_digest) pairs. :param hexdigests_remote: The expected hexdigests as (algorithm_name, url to file with expected hexdigest) pairs. :param hexdigests_strict: Set this to ``False`` to stop automatically checking for the `algorithm(filename)=hash` format :param progress_bar: Set to true to show a progress bar while downloading :param tqdm_kwargs: Override the default arguments passed to :class:`tadm.tqdm` when progress_bar is True. :param kwargs: If using :func:`urllib.request.urlretrieve`, there are no keyword arguments available. If using ``requests`` as a backend, passes these to :func:`requests.get`. If using ``requests`` as a backend, ``stream`` is set to True by default. :raises Exception: Thrown if an error besides a keyboard interrupt is thrown during download :raises KeyboardInterrupt: If a keyboard interrupt is thrown during download :raises UnexpectedDirectory: If a directory is given for the ``path`` argument :raises ValueError: If an invalid backend is chosen :raises DownloadError: If an error occurs during download """ path = Path(path).resolve() if path.is_dir(): raise UnexpectedDirectoryError(path) if path.is_file() and not force: raise_on_digest_mismatch( path=path, hexdigests=hexdigests, hexdigests_remote=hexdigests_remote, hexdigests_strict=hexdigests_strict, ) logger.debug("did not re-download %s from %s", path, url) return if backend is None: backend = "urllib" desc = f"Downloading {path.name}" if _version: desc += f" (v{_version})" _tqdm_kwargs = { "unit": "B", "unit_scale": True, "unit_divisor": 1024, "miniters": 1, "disable": not progress_bar, "desc": desc, "leave": False, } if tqdm_kwargs: _tqdm_kwargs.update(tqdm_kwargs) try: if backend == "urllib": logger.info("downloading with urllib from %s to %s", url, path) if kwargs: logger.warning( "no kwargs should be supplied when using urllib, skipping: %s", kwargs ) with TqdmReportHook(**_tqdm_kwargs) as t: try: urlretrieve(url, path, reporthook=t.update_to) # noqa:S310 except urllib.error.URLError as e: logger.info("download failed %s to %s", url, e) raise DownloadError(backend, url, path, e) from e elif backend == "requests": import requests import requests.exceptions kwargs.setdefault("stream", True) if "headers" not in kwargs or kwargs["headers"] is None: kwargs["headers"] = {} # ignore the type error because we make sure it's a dict above kwargs["headers"].setdefault("User-Agent", DEFAULT_AGENT) # type:ignore[union-attr] try: # see https://requests.readthedocs.io/en/master/user/quickstart/#raw-response-content # pattern from https://stackoverflow.com/a/39217788/5775947 with requests.get(url, **kwargs) as response, path.open("wb") as file: # noqa:S113 logger.info( "downloading (stream=%s) with requests from %s to %s", kwargs["stream"], url, path, ) # Solution for progress bar from https://stackoverflow.com/a/63831344/5775947 total_size = int(response.headers.get("Content-Length", 0)) # Decompress if needed response.raw.read = partial( # type:ignore[method-assign] response.raw.read, decode_content=True ) with tqdm.wrapattr( response.raw, "read", total=total_size, **_tqdm_kwargs ) as fsrc: shutil.copyfileobj(fsrc, file) except requests.exceptions.ConnectionError as e: raise DownloadError(backend, url, path, e) from e else: raise ValueError(f'Invalid backend: {backend}. Use "requests" or "urllib".') except (Exception, KeyboardInterrupt): if clean_on_failure: path.unlink(missing_ok=True) raise raise_on_digest_mismatch( path=path, hexdigests=hexdigests, hexdigests_remote=hexdigests_remote, hexdigests_strict=hexdigests_strict, ) class DownloadError(OSError): """An error that wraps information from a requests or urllib download failure.""" def __init__( self, backend: DownloadBackend, url: str, path: Path, exc: urllib.error.URLError | requests.exceptions.ConnectionError, ) -> None: """Initialize the error. :param backend: The backend used :param url: The url that failed to download :param path: The path that was supposed to be downloaded to :param exc: The exception raised """ self.backend = backend self.url = url self.path = path self.exc = exc # TODO parse out HTTP error code, if possible def __str__(self) -> str: return f"Failed with {self.backend} to download {self.url} to {self.path}" class UnexpectedDirectoryError(FileExistsError): """Thrown if a directory path is given where file path should have been.""" def __init__(self, path: Path): """Instantiate the exception. :param path: The path to a directory that should have been a file. """ self.path = path def __str__(self) -> str: return f"got directory instead of file: {self.path}" CHUNK_SIZE = 32768 DOWNLOAD_URL = "https://docs.google.com/uc?export=download" TOKEN_KEY = "download_warning" # noqa:S105 def download_from_google( file_id: str, path: str | Path, force: bool = True, clean_on_failure: bool = True, hexdigests: Mapping[str, str] | None = None, ) -> None: """Download a file from google drive. Implementation inspired by https://github.com/ndrplz/google-drive-downloader. :param file_id: The google file identifier :param path: The place to write the file :param force: If false and the file already exists, will not re-download. :param clean_on_failure: If true, will delete the file on any exception raised during download :param hexdigests: The expected hexdigests as (algorithm_name, expected_hex_digest) pairs. :raises Exception: Thrown if an error besides a keyboard interrupt is thrown during download :raises KeyboardInterrupt: If a keyboard interrupt is thrown during download :raises UnexpectedDirectory: If a directory is given for the ``path`` argument """ import requests path = Path(path).resolve() if path.is_dir(): raise UnexpectedDirectoryError(path) if path.is_file() and not force: raise_on_digest_mismatch(path=path, hexdigests=hexdigests) logger.debug("did not re-download %s from Google ID %s", path, file_id) return try: with requests.Session() as sess: res = sess.get(DOWNLOAD_URL, params={"id": file_id}, stream=True) token = _get_confirm_token(res) res = sess.get(DOWNLOAD_URL, params={"id": file_id, "confirm": token}, stream=True) with path.open("wb") as file: for chunk in tqdm(res.iter_content(CHUNK_SIZE), desc="writing", unit="chunk"): if chunk: # filter out keep-alive new chunks file.write(chunk) except (Exception, KeyboardInterrupt): if clean_on_failure: path.unlink(missing_ok=True) raise raise_on_digest_mismatch(path=path, hexdigests=hexdigests) def _get_confirm_token(res: requests.Response) -> str: for key, value in res.cookies.items(): if key.startswith(TOKEN_KEY): return value raise ValueError(f"no token found with key {TOKEN_KEY} in cookies: {res.cookies}") def download_from_s3( s3_bucket: str, s3_key: str, path: str | Path, client: botocore.client.BaseClient | None = None, client_kwargs: Mapping[str, Any] | None = None, download_file_kwargs: Mapping[str, Any] | None = None, force: bool = True, clean_on_failure: bool = True, ) -> None: """Download a file from S3. :param s3_bucket: The key inside the S3 bucket name :param s3_key: The key inside the S3 bucket :param path: The place to write the file :param client: A botocore client. If none given, one will be created automatically :param client_kwargs: Keyword arguments to be passed to the client on instantiation. :param download_file_kwargs: Keyword arguments to be passed to :func:`boto3.s3.transfer.S3Transfer.download_file` :param force: If false and the file already exists, will not re-download. :param clean_on_failure: If true, will delete the file on any exception raised during download :raises Exception: Thrown if an error besides a keyboard interrupt is thrown during download :raises KeyboardInterrupt: If a keyboard interrupt is thrown during download :raises UnexpectedDirectory: If a directory is given for the ``path`` argument """ path = Path(path).resolve() if path.is_dir(): raise UnexpectedDirectoryError(path) if path.is_file() and not force: logger.debug("did not re-download %s from %s %s", path, s3_bucket, s3_key) return try: import boto3.s3.transfer if client is None: import boto3 import botocore.client client_kwargs = {} if client_kwargs is None else dict(client_kwargs) client_kwargs.setdefault( "config", botocore.client.Config(signature_version=botocore.UNSIGNED) ) client = boto3.client("s3", **client_kwargs) download_file_kwargs = {} if download_file_kwargs is None else dict(download_file_kwargs) download_file_kwargs.setdefault( "Config", boto3.s3.transfer.TransferConfig(use_threads=False) ) client.download_file(s3_bucket, s3_key, path.as_posix(), **download_file_kwargs) except (Exception, KeyboardInterrupt): if clean_on_failure: path.unlink(missing_ok=True) raise cthoyt-pystow-54a586e/src/pystow/utils/env.py000066400000000000000000000115141525577235500214120ustar00rootroot00000000000000"""Environment utilities.""" from __future__ import annotations import contextlib import os import tempfile from collections.abc import Iterator from pathlib import Path from ..constants import ( PYSTOW_HOME_ENVVAR, PYSTOW_NAME_DEFAULT, PYSTOW_NAME_ENVVAR, PYSTOW_USE_APPDIRS, ) __all__ = [ "get_base", "get_home", "get_name", "getenv_path", "mkdir", "mock_envvar", "mock_home", "use_appdirs", ] def mkdir(path: Path, ensure_exists: bool = True) -> None: """Make a directory (or parent directory if a file is given) if flagged with ``ensure_exists``. :param path: The path to a directory :param ensure_exists: Should the directories leading to the path be created if they don't already exist? """ if ensure_exists: path.mkdir(exist_ok=True, parents=True) @contextlib.contextmanager def mock_envvar(envvar: str, value: str) -> Iterator[None]: """Mock the environment variable then delete it after the test is over. :param envvar: The environment variable to mock :param value: The value to temporarily put in the environment variable during this mock. :yield: None, since this just mocks the environment variable for the time being. """ original_value = os.environ.get(envvar) os.environ[envvar] = value yield if original_value is None: del os.environ[envvar] else: os.environ[envvar] = original_value @contextlib.contextmanager def mock_home() -> Iterator[Path]: """Mock the PyStow home environment variable, yields the directory name. :yield: The path to the temporary directory. """ with tempfile.TemporaryDirectory() as directory, mock_envvar(PYSTOW_HOME_ENVVAR, directory): yield Path(directory) def getenv_path(envvar: str, default: Path, ensure_exists: bool = True) -> Path: """Get an environment variable representing a path, or use the default. :param envvar: The environmental variable name to check :param default: The default path to return if the environmental variable is not set :param ensure_exists: Should the directories leading to the path be created if they don't already exist? :returns: A path either specified by the environmental variable or by the default. """ rv = Path(os.getenv(envvar, default=default)).expanduser() mkdir(rv, ensure_exists=ensure_exists) return rv def use_appdirs() -> bool: """Check if X Desktop Group (XDG) compatibility is requested. :returns: If the :data:`PYSTOW_USE_APPDIRS` is set to ``true`` in the environment. """ return os.getenv(PYSTOW_USE_APPDIRS) in {"true", "True"} def get_home(ensure_exists: bool = True) -> Path: """Get the PyStow home directory. :param ensure_exists: If true, ensures the directory is created :returns: A path object representing the pystow home directory, as one of: 1. :data:`PYSTOW_HOME_ENVVAR` environment variable or 2. The user data directory defined by :mod:`appdirs` or :mod:`platformdirs` if the :data:`PYSTOW_USE_APPDIRS` environment variable is set to ``true`` or 3. The default directory constructed in the user's home directory plus what's returned by :func:`get_name`. """ if use_appdirs(): try: from platformdirs import user_data_dir except ImportError: from appdirs import user_data_dir default = Path(user_data_dir()) else: default = Path.home() / get_name() return getenv_path(PYSTOW_HOME_ENVVAR, default, ensure_exists=ensure_exists) def get_base(key: str, ensure_exists: bool = True) -> Path: """Get the base directory for a module. :param key: The name of the module. No funny characters. The envvar _HOME where key is uppercased is checked first before using the default home directory. :param ensure_exists: Should all directories be created automatically? Defaults to true. :returns: The path to the given :raises ValueError: if the key is invalid (e.g., has a dot in it) """ if "." in key: raise ValueError(f"The module should not have a dot in it: {key}") envvar = f"{key.upper()}_HOME" if use_appdirs(): try: from platformdirs import user_data_dir except ImportError: from appdirs import user_data_dir default = Path(user_data_dir(appname=key)) else: default = get_home(ensure_exists=False) / key return getenv_path(envvar, default, ensure_exists=ensure_exists) def get_name() -> str: """Get the PyStow home directory name. :returns: The name of the pystow home directory, either loaded from the :data:`PYSTOW_NAME_ENVVAR` environment variable or given by the default value :data:`PYSTOW_NAME_DEFAULT`. """ return os.getenv(PYSTOW_NAME_ENVVAR, default=PYSTOW_NAME_DEFAULT) cthoyt-pystow-54a586e/src/pystow/utils/hashing.py000066400000000000000000000147711525577235500222530ustar00rootroot00000000000000"""Hashing utilities.""" from __future__ import annotations import hashlib import logging from collections.abc import Collection, Iterable, Mapping from pathlib import Path from typing import NamedTuple, TypeAlias from urllib.request import urlopen __all__ = [ "Hash", "HexDigestError", "HexDigestMismatch", "get_hash_hexdigest", "get_hashes", "get_hexdigests_remote", "get_offending_hexdigests", "raise_on_digest_mismatch", ] logger = logging.getLogger(__name__) #: This type alias uses a stub-only constructor, meaning that #: hashlib._Hash isn't actually part of the code, but MyPy injects it #: so we can do type checking Hash: TypeAlias = "hashlib._Hash" def get_hexdigests_remote( hexdigests_remote: Mapping[str, str] | None, hexdigests_strict: bool = False ) -> Mapping[str, str]: """Process hexdigests via URLs. :param hexdigests_remote: The expected hexdigests as (algorithm_name, url to file with expected hex digest) pairs. :param hexdigests_strict: Set this to `False` to stop automatically checking for the `algorithm(filename)=hash` format :returns: A mapping of algorithms to hexdigests """ rv = {} for key, url in (hexdigests_remote or {}).items(): with urlopen(url) as response: # noqa:S310 text = response.read().decode("utf-8") if not hexdigests_strict and "=" in text: text = text.rsplit("=", 1)[-1].strip() rv[key] = text return rv class HexDigestMismatch(NamedTuple): """Contains information about a hexdigest mismatch.""" #: the name of the algorithm name: str #: the observed/actual hexdigest, encoded as a string actual: str #: the expected hexdigest, encoded as a string expected: str def get_offending_hexdigests( path: str | Path, chunk_size: int | None = None, hexdigests: Mapping[str, str] | None = None, hexdigests_remote: Mapping[str, str] | None = None, hexdigests_strict: bool = False, ) -> Collection[HexDigestMismatch]: """Check a file for hash sums. :param path: The file path. :param chunk_size: The chunk size for reading the file. :param hexdigests: The expected hexdigests as (algorithm_name, expected_hex_digest) pairs. :param hexdigests_remote: The expected hexdigests as (algorithm_name, url to file with expected hexdigest) pairs. :param hexdigests_strict: Set this to false to stop automatically checking for the `algorithm(filename)=hash` format :returns: A collection of observed / expected hexdigests where the digests do not match. """ hexdigests = dict( **(hexdigests or {}), **get_hexdigests_remote(hexdigests_remote, hexdigests_strict=hexdigests_strict), ) # If there aren't any keys in the combine dictionaries, # then there won't be any mismatches if not hexdigests: return [] logger.info(f"Checking hash sums for file: {path}") # instantiate algorithms algorithms = get_hashes(path=path, names=set(hexdigests), chunk_size=chunk_size) # Compare digests mismatches = [] for alg, expected_digest in hexdigests.items(): observed_digest = algorithms[alg].hexdigest() if observed_digest != expected_digest: logger.error(f"{alg} expected {expected_digest} but got {observed_digest}.") mismatches.append(HexDigestMismatch(alg, observed_digest, expected_digest)) else: logger.debug(f"Successfully checked with {alg}.") return mismatches def get_hashes( path: str | Path, names: Iterable[str], *, chunk_size: int | None = None, ) -> Mapping[str, Hash]: """Calculate several hexdigests of hash algorithms for a file concurrently. :param path: The file path. :param names: Names of the hash algorithms in :mod:`hashlib` :param chunk_size: The chunk size for reading the file. :returns: A collection of observed hexdigests """ path = Path(path).resolve() if chunk_size is None: chunk_size = 64 * 2**10 # instantiate hash algorithms algorithms: Mapping[str, Hash] = {name: hashlib.new(name) for name in names} # calculate hash sums of file incrementally buffer = memoryview(bytearray(chunk_size)) with path.open("rb", buffering=0) as file: for this_chunk_size in iter(lambda: file.readinto(buffer), 0): for alg in algorithms.values(): alg.update(buffer[:this_chunk_size]) return algorithms def get_hash_hexdigest( path: str | Path, name: str, *, chunk_size: int | None = None, ) -> str: """Get a hash digest for a single hash.""" r = get_hashes(path, [name], chunk_size=chunk_size) return r[name].hexdigest() def raise_on_digest_mismatch( *, path: Path, hexdigests: Mapping[str, str] | None = None, hexdigests_remote: Mapping[str, str] | None = None, hexdigests_strict: bool = False, ) -> None: """Raise a HexDigestError if the digests do not match. :param path: The file path. :param hexdigests: The expected hexdigests as (algorithm_name, expected_hex_digest) pairs. :param hexdigests_remote: The expected hexdigests as (algorithm_name, url to file with expected hexdigest) pairs. :param hexdigests_strict: Set this to false to stop automatically checking for the `algorithm(filename)=hash` format :raises HexDigestError: if there are any offending hex digests The expected hexdigests as (algorithm_name, url to file with expected hexdigest) pairs. """ offending_hexdigests = get_offending_hexdigests( path=path, hexdigests=hexdigests, hexdigests_remote=hexdigests_remote, hexdigests_strict=hexdigests_strict, ) if offending_hexdigests: raise HexDigestError(offending_hexdigests) class HexDigestError(ValueError): """Thrown if the hashsums do not match expected hashsums.""" def __init__(self, offending_hexdigests: Collection[HexDigestMismatch]): """Instantiate the exception. :param offending_hexdigests: The result from :func:`get_offending_hexdigests` """ self.offending_hexdigests = offending_hexdigests def __str__(self) -> str: return "\n".join( ( "Hexdigest of downloaded file does not match the expected ones!", *( f"\t{name} actual: {actual} vs. expected: {expected}" for name, actual, expected in self.offending_hexdigests ), ) ) cthoyt-pystow-54a586e/src/pystow/utils/io_typing.py000066400000000000000000000115171525577235500226260ustar00rootroot00000000000000"""Typing for I/O.""" from __future__ import annotations import typing from collections.abc import Mapping from typing import TYPE_CHECKING, Literal, TypeAlias if TYPE_CHECKING: import _csv __all__ = [ "MODE_MAP", "OPERATION_VALUES", "REPRESENTATION_VALUES", "REVERSE_MODE_MAP", "_OPERATION_TO_UNQUALIFIED_MODE", "InvalidOperationError", "InvalidRepresentationError", "Operation", "Reader", "Representation", "Writer", "ensure_sensible_default_encoding", "ensure_sensible_newline", "get_mode_pair", ] Reader: TypeAlias = "_csv.Reader" Writer: TypeAlias = "_csv.Writer" #: A human-readable flag for how to open a file. Operation: TypeAlias = Literal["read", "write"] OPERATION_VALUES: set[str] = set(typing.get_args(Operation)) #: A human-readable flag for how to open a file. Representation: TypeAlias = Literal["text", "binary"] REPRESENTATION_VALUES: set[str] = set(typing.get_args(Representation)) #: Characters for "unqualified" modes, which might be interpreted #: differently by different functions UnqualifiedMode: TypeAlias = Literal["r", "w"] #: Characters for "qualified" modes, which are absolute (as opposed to #: :data:`UnqualifiedMode`, which is context-dependent) QualifiedMode: TypeAlias = Literal["rt", "wt", "rb", "wb"] ModePair: TypeAlias = tuple[Operation, Representation] _OPERATION_TO_UNQUALIFIED_MODE: Mapping[Operation, UnqualifiedMode] = { "read": "r", "write": "w", } _OPERATION_TO_BINARY_MODE: Mapping[Operation, Literal["rb", "wb"]] = { "read": "rb", "write": "wb", } #: A mapping between operation/representation pairs and qualified modes MODE_MAP: dict[ModePair, QualifiedMode] = { ("read", "text"): "rt", ("read", "binary"): "rb", ("write", "text"): "wt", ("write", "binary"): "wb", } #: A mapping between qualified modes and operation/representation pairs REVERSE_MODE_MAP: dict[QualifiedMode, ModePair] = { "rt": ("read", "text"), "rb": ("read", "binary"), "wt": ("write", "text"), "wb": ("write", "binary"), } UNQUALIFIED_TEXT_MAP: dict[UnqualifiedMode, ModePair] = { "r": ("read", "text"), "w": ("write", "text"), } UNQUALIFIED_BINARY_MAP: dict[UnqualifiedMode, ModePair] = { "r": ("read", "binary"), "w": ("write", "binary"), } def get_mode_pair( mode: UnqualifiedMode | QualifiedMode, interpretation: Representation ) -> ModePair: """Get the mode pair.""" match mode: case "rt" | "wt" | "rb" | "wb": return REVERSE_MODE_MAP[mode] case "r" | "w" if interpretation == "text": return UNQUALIFIED_TEXT_MAP[mode] case "r" | "w" if interpretation == "binary": return UNQUALIFIED_BINARY_MAP[mode] case _: raise ValueError(f"invalid mode: {mode}") class InvalidRepresentationError(ValueError): """Raised when passing an invalid representation.""" def __init__(self, representation: str) -> None: """Instantiate the exception.""" self.representation = representation def __str__(self) -> str: """Create a string for the exception.""" return ( f"Invalid representation: {self.representation}. " f"Should be one of {REPRESENTATION_VALUES}." ) class InvalidOperationError(ValueError): """Raised when passing an invalid operation.""" def __init__(self, operation: str) -> None: """Instantiate the exception.""" self.operation = operation def __str__(self) -> str: """Create a string for the exception.""" return f"Invalid operation: {self.operation}. Should be one of {OPERATION_VALUES}." def ensure_sensible_default_encoding( encoding: str | None, *, representation: Representation ) -> str | None: """Get a sensible default encoding.""" # this function exists because windows doesn't use UTF-8 as a default # encoding for some reason, and that's bonk. So we intercept the encoding # and set it explicitly to UTF-8 if representation == "binary": if encoding is not None: raise ValueError else: return None elif representation == "text": if encoding is not None: return encoding return "utf-8" else: raise InvalidRepresentationError(representation) def ensure_sensible_newline(newline: str | None, *, representation: Representation) -> str | None: """Get a sensible default newline.""" # this function exists to override the default way newlines are # automatically interpreted by python on Windows to always use # \n instead of \r\n if representation == "binary": if newline is not None: raise ValueError return None elif representation == "text": if newline is not None: return newline return "\n" else: raise InvalidRepresentationError(representation) cthoyt-pystow-54a586e/src/pystow/utils/iter.py000066400000000000000000000064001525577235500215630ustar00rootroot00000000000000"""Utilities for functions that consume iterables.""" from __future__ import annotations from typing import TYPE_CHECKING, Concatenate, ParamSpec, TypeVar if TYPE_CHECKING: from collections.abc import Callable, Generator, Iterable X = TypeVar("X") P = ParamSpec("P") __all__ = [ "reyield", ] class Sentinel: """A sentinel class.""" _SENTINEL = Sentinel() def reyield( func: Callable[Concatenate[Iterable[X], P], None], elements: Iterable[X], *args: P.args, **kwargs: P.kwargs, ) -> Generator[X]: """Make a function that consumes an iterable yield its elements. :param func: A function that consumes an iterable, and does not return anything :param elements: An iterable :returns: A generator that injects the elements of the iterable one at a time before yielding them. >>> state = {"value": 0} >>> def accumulate(elements: Iterable[int]) -> None: ... for element in elements: ... state["value"] += element >>> square_sum = 0 >>> for element in reyield(accumulate, range(5)): ... square_sum += element**2 The goal of this function is not to need to keep elements in memory, so this should be an iterable-native alternative to just doing: >>> elements = range(5) >>> elements = list(elements) >>> accumulate(elements) >>> square_sum = 0 >>> for element in elements: ... square_sum += element**2 """ from contextlib import closing with closing(_help_reiter(func, *args, **kwargs)) as generator: next(generator) # prime the coroutine for element in elements: generator.send(element) yield element def _help_reiter( func: Callable[Concatenate[Iterable[X], P], None], *args: P.args, **kwargs: P.kwargs ) -> Generator[None, X, None]: """Modify the function.""" import threading from queue import Queue queue: Queue[X | Sentinel] = Queue() def _iterable_from_queue() -> Generator[X]: while True: # queue.get() blocks indefinitely (because timeout is None) # until the queue receives something item_ = queue.get(block=True, timeout=None) # if the queue receives a sentinel value, then we break out of the # while loop, which will cause the generator to raise a GeneratorExit if isinstance(item_, Sentinel): break yield item_ # Run the consumer in a background thread, fed by the generator that wraps the queue thread = threading.Thread( target=func, args=(_iterable_from_queue(), *args), kwargs=kwargs, daemon=True ) thread.start() # now, we invert the generator pattern - this function # will return a generator that you can .send(item) to, # and the yield statement waits for the value and sticks # it in item try: while True: item: X = yield # pause, wait for .send(item) queue.put(item) except GeneratorExit: # this happens when the outer generator ends, which means # it's time to send the queue a "final" value and block # on finishing the thread to clean up pass finally: queue.put(_SENTINEL) # signal the consumer that we're done thread.join() cthoyt-pystow-54a586e/src/pystow/utils/pydantic_utils.py000066400000000000000000000166351525577235500236660ustar00rootroot00000000000000"""Utilities for working with Pydantic.""" from __future__ import annotations import logging import typing from collections.abc import Callable, Generator, Iterable, Mapping from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, TextIO, TypeAlias from tqdm import tqdm from .safe_open import ( _open_read_text, _open_write_text, safe_open_dict_reader, safe_open_json, safe_open_yaml, write_json, write_yaml, ) if TYPE_CHECKING: import pydantic __all__ = [ "ModelValidateFailureAction", "iter_pydantic_jsonl", "iter_pydantic_tsv", "model_dump_yaml", "read_pydantic_json", "read_pydantic_jsonl", "read_pydantic_tsv", "read_pydantic_yaml", "stream_write_pydantic_jsonl", "write_pydantic_json", "write_pydantic_jsonl", "write_pydantic_yaml", ] logger = logging.getLogger(__name__) BaseModelVar = typing.TypeVar("BaseModelVar", bound="pydantic.BaseModel") #: The action to take on model validation failure ModelValidateFailureAction: TypeAlias = Literal["raise", "skip"] def iter_pydantic_jsonl( file: str | Path | TextIO, model_cls: type[BaseModelVar], *, progress: bool = False, failure_action: ModelValidateFailureAction = "skip", encoding: str | None = None, newline: str | None = None, tqdm_kwargs: Mapping[str, Any] | None = None, ) -> Iterable[BaseModelVar]: """Read models to a file as JSONL.""" import pydantic _tqdm_kwargs = { "desc": "Reading mappings", "leave": False, "unit": "mapping", "unit_scale": True, } if tqdm_kwargs is not None: _tqdm_kwargs.update(tqdm_kwargs) with _open_read_text(file, encoding=encoding, newline=newline) as file_: for i, line in enumerate(tqdm(file_, disable=not progress, **_tqdm_kwargs)): try: yv = model_cls.model_validate_json(line.strip()) except pydantic.ValidationError: if failure_action == "raise": raise else: logger.debug("[line:%d] failed to parse JSON", i) continue else: yield yv def read_pydantic_jsonl( file: str | Path | TextIO, model_cls: type[BaseModelVar], **kwargs: Any ) -> list[BaseModelVar]: """Read models from a file as JSONL.""" return list(iter_pydantic_jsonl(file, model_cls, **kwargs)) def write_pydantic_jsonl( models: Iterable[pydantic.BaseModel], file: str | Path | TextIO, **kwargs: Any ) -> None: """Write models to a file as JSONL.""" kwargs.setdefault("exclude_none", True) kwargs.setdefault("exclude_unset", True) kwargs.setdefault("exclude_defaults", True) with _open_write_text(file) as file_: for model in models: file_.write(model.model_dump_json(**kwargs) + "\n") def stream_write_pydantic_jsonl( models: Iterable[BaseModelVar], file: str | Path | TextIO, **kwargs: Any ) -> Generator[BaseModelVar]: """Write models to a file as JSONL and yield them.""" kwargs.setdefault("exclude_none", True) kwargs.setdefault("exclude_unset", True) kwargs.setdefault("exclude_defaults", True) with _open_write_text(file) as file_: for model in models: file_.write(model.model_dump_json(**kwargs) + "\n") yield model def read_pydantic_tsv( path: str | Path | TextIO, model: type[BaseModelVar], *, process: Callable[[dict[str, Any]], dict[str, Any]] | None = None, failure_action: ModelValidateFailureAction = "skip", ) -> list[BaseModelVar]: """Read models from a TSV file.""" return list(iter_pydantic_tsv(path, model, process=process, failure_action=failure_action)) def iter_pydantic_tsv( path: str | Path | TextIO, model_cls: type[BaseModelVar], *, process: Callable[[dict[str, Any]], dict[str, Any]] | None = None, failure_action: ModelValidateFailureAction = "skip", ) -> Generator[BaseModelVar]: """Read models from a TSV file, iteratively.""" import pydantic with safe_open_dict_reader(path) as reader: records: Iterable[dict[str, Any]] if process is None: records = iter(reader) else: records = (process(record) for record in reader) for record in records: try: yv = model_cls.model_validate(record) except pydantic.ValidationError: if failure_action == "raise": raise else: logger.debug("[line:%d] failed to parse row", record) continue else: yield yv def read_pydantic_json( path_or_url: str | Path | TextIO, model_cls: type[BaseModelVar], *, encoding: str | None = None, newline: str | None = None, ) -> BaseModelVar: """Read a JSON file into a model.""" return model_cls.model_validate(safe_open_json(path_or_url, encoding=encoding, newline=newline)) def read_pydantic_yaml( path_or_url: str | Path | TextIO, model_cls: type[BaseModelVar], *, encoding: str | None = None, newline: str | None = None, ) -> BaseModelVar: """Read a YAML file into a model.""" return model_cls.model_validate(safe_open_yaml(path_or_url, encoding=encoding, newline=newline)) def model_dump_yaml( model: pydantic.BaseModel, *, exclude_none: bool = False, exclude_unset: bool = False, exclude: set[str] | None = None, allow_unicode: bool = True, indent: int | None = None, sort_keys: bool = False, ) -> str: """Dump the model as YAML string.""" import yaml data = model.model_dump( mode="json", exclude_none=exclude_none, exclude_unset=exclude_unset, exclude=exclude ) return yaml.safe_dump(data, allow_unicode=allow_unicode, indent=indent, sort_keys=sort_keys) def write_pydantic_yaml( model: pydantic.BaseModel, path: str | Path | TextIO, *, exclude: set[str] | None = None, exclude_none: bool = True, exclude_unset: bool = True, exclude_defaults: bool = True, encoding: str | None = None, newline: str | None = None, indent: int | None = None, allow_unicode: bool = True, sort_keys: bool = False, ) -> None: """Write a model to a YAML file.""" data = model.model_dump( mode="json", exclude_none=exclude_none, exclude_unset=exclude_unset, exclude_defaults=exclude_defaults, exclude=exclude, ) write_yaml( data, path, encoding=encoding, newline=newline, indent=indent, allow_unicode=allow_unicode, sort_keys=sort_keys, ) def write_pydantic_json( model: pydantic.BaseModel, path: str | Path | TextIO, *, exclude: set[str] | None = None, exclude_none: bool = True, exclude_unset: bool = True, exclude_defaults: bool = True, encoding: str | None = None, newline: str | None = None, ensure_ascii: bool = False, indent: int | None = None, trailing_newline: bool = True, ) -> None: """Write a model to a JSON file.""" data = model.model_dump( mode="json", exclude=exclude, exclude_none=exclude_none, exclude_unset=exclude_unset, exclude_defaults=exclude_defaults, ) write_json( data, path, encoding=encoding, newline=newline, ensure_ascii=ensure_ascii, indent=indent, trailing_newline=trailing_newline, ) cthoyt-pystow-54a586e/src/pystow/utils/safe_open.py000066400000000000000000000320051525577235500225570ustar00rootroot00000000000000"""File opening utilities.""" from __future__ import annotations import bz2 import contextlib import csv import gzip import io import json import lzma import sys import typing import urllib.request import zipfile from collections.abc import Generator, Mapping from pathlib import Path, PurePath, PurePosixPath from typing import ( IO, Any, Literal, Never, NotRequired, TypeAlias, TypedDict, TypeGuard, Unpack, overload, ) from .io_typing import ( _OPERATION_TO_BINARY_MODE, _OPERATION_TO_UNQUALIFIED_MODE, MODE_MAP, OPERATION_VALUES, REPRESENTATION_VALUES, InvalidOperationError, InvalidRepresentationError, Operation, Representation, ensure_sensible_default_encoding, ensure_sensible_newline, ) try: if sys.version_info >= (3, 14): from compression import zstd else: from backports import zstd zstd_open = zstd.open except ImportError: def zstd_open(*args: Any, **kwargs: Any) -> Never: """Open a Zstandard compressed file in binary or text mode.""" raise RuntimeError("zstd is not available") zstd_available = False else: zstd_available = True __all__ = [ "is_url", "open_inner_zipfile", "open_url", "safe_open", "safe_open_dict_reader", "safe_open_json", "safe_open_yaml", "safe_read_text", "safe_write_text", "write_json", "write_yaml", "zstd_open", ] TextSource: TypeAlias = str | Path | IO[str] COMPRESSION_EXTENSIONS = ["gz", "xz", "bz2"] if zstd_available: COMPRESSION_EXTENSIONS.append("zst") class OpenKwargs(TypedDict): """Keyword arguments for open-like functions.""" encoding: NotRequired[str | None] newline: NotRequired[str | None] buffering: NotRequired[int | None] timeout: NotRequired[int | None] """Timeout only applies when opening from a remote file""" # docstr-coverage:excused `overload` @typing.overload @contextlib.contextmanager def safe_open( path: str | Path | IO[str] | IO[bytes], *, operation: Operation = ..., representation: Literal["text"] = "text", **kwargs: Unpack[OpenKwargs], ) -> Generator[IO[str]]: ... # docstr-coverage:excused `overload` @typing.overload @contextlib.contextmanager def safe_open( path: str | Path | IO[str] | IO[bytes], *, operation: Operation = ..., representation: Literal["binary"] = "binary", **kwargs: Unpack[OpenKwargs], ) -> Generator[IO[bytes]]: ... @contextlib.contextmanager def safe_open( # noqa:C901 path: str | Path | IO[str] | IO[bytes], *, operation: Operation = "read", representation: Representation = "text", **kwargs: Unpack[OpenKwargs], ) -> Generator[IO[str]] | Generator[IO[bytes]]: """Safely open a file for reading or writing text.""" if operation not in OPERATION_VALUES: raise InvalidOperationError(operation) if representation not in REPRESENTATION_VALUES: raise InvalidRepresentationError(representation) if isinstance(path, (str, Path)): encoding = ensure_sensible_default_encoding( kwargs.get("encoding"), representation=representation ) newline = ensure_sensible_newline(kwargs.get("newline"), representation=representation) buffering = kwargs.get("buffering") or -1 if is_url(path): if operation != "read": raise ValueError('can only use operation="read" with URLs') with open_url( path, representation=representation, encoding=encoding, newline=newline, buffering=buffering, timeout=kwargs.get("timeout"), ) as file: yield file else: mode = MODE_MAP[operation, representation] premode = _OPERATION_TO_BINARY_MODE[operation] path = Path(path).expanduser().resolve() if path.suffix.endswith(".gz"): with ( open(path, buffering=buffering, mode=premode) as raw, gzip.open(raw, mode=mode, encoding=encoding, newline=newline) as gzf, ): yield gzf # type:ignore elif path.suffix.endswith(".bz2"): with ( open(path, buffering=buffering, mode=premode) as raw, bz2.open(raw, mode=mode, encoding=encoding, newline=newline) as bz2f, ): yield bz2f elif path.suffix.endswith(".xz"): with ( open(path, buffering=buffering, mode=premode) as raw, lzma.open(raw, mode=mode, encoding=encoding, newline=newline) as lzmaf, ): yield lzmaf elif path.suffix.endswith(".zst"): with ( open(path, buffering=buffering, mode=premode) as raw, zstd_open(raw, mode=mode, encoding=encoding, newline=newline) as zstdf, ): yield zstdf # type:ignore else: with open( path, mode=mode, encoding=encoding, newline=newline, buffering=buffering ) as file: yield file elif isinstance(path, typing.TextIO | io.TextIOWrapper | io.TextIOBase): if representation != "text": path = path.buffer yield path # io.BufferedIOBase covers the LZMA, BZ2, Gzip, and ZSTD file types # as well as io.BufferedReader elif isinstance(path, typing.BinaryIO | io.BufferedIOBase): with _wrap_binary_if_needed( path, representation, encoding=kwargs.get("encoding"), newline=kwargs.get("newline") ) as yp: yield yp else: raise TypeError(f"unsupported type for opening: {type(path)} - {path}") @contextlib.contextmanager def _open_read_text(path: TextSource, **kwargs: Unpack[OpenKwargs]) -> Generator[IO[str]]: with safe_open(path, representation="text", operation="read", **kwargs) as file: yield file @contextlib.contextmanager def _open_write_text(path: TextSource, **kwargs: Unpack[OpenKwargs]) -> Generator[IO[str]]: with safe_open(path, representation="text", operation="write", **kwargs) as file: yield file def safe_open_json(path_or_url: TextSource, **kwargs: Unpack[OpenKwargs]) -> Any: """Safely open a file and parse as JSON.""" with _open_read_text(path_or_url, **kwargs) as file: return json.load(file) def safe_open_yaml(path_or_url: TextSource, **kwargs: Unpack[OpenKwargs]) -> Any: """Safely open a file and parse as YAML.""" import yaml with _open_read_text(path_or_url, **kwargs) as file: return yaml.safe_load(file) def safe_write_text(s: str, path: TextSource, **kwargs: Unpack[OpenKwargs]) -> int: """Write text to a file.""" with _open_write_text(path, **kwargs) as file: return file.write(s) def safe_read_text(path: TextSource, **kwargs: Unpack[OpenKwargs]) -> str: """Read text from a file.""" with _open_read_text(path, **kwargs) as file: return file.read() def write_yaml( data: Any, path: TextSource, *, encoding: str | None = None, newline: str | None = None, buffering: int | None = None, indent: int | None = None, allow_unicode: bool = True, **kwargs: Any, ) -> Any: """Write YAML to a file.""" import yaml with _open_write_text(path, encoding=encoding, newline=newline, buffering=buffering) as file: yaml.safe_dump(data, file, indent=indent, allow_unicode=allow_unicode, **kwargs) def write_json( data: Any, path: TextSource, *, encoding: str | None = None, newline: str | None = None, ensure_ascii: bool = False, indent: int | None = None, trailing_newline: bool = True, **kwargs: Any, ) -> Any: """Write JSON to a file.""" with _open_write_text(path, encoding=encoding, newline=newline) as file: json.dump(data, file, ensure_ascii=ensure_ascii, indent=indent, **kwargs) if trailing_newline: file.write("\n") # docstr-coverage:excused `overload` @typing.overload @contextlib.contextmanager def open_inner_zipfile( zip_file: zipfile.ZipFile, inner_path: str | PurePath, *, operation: Operation = ..., representation: Literal["text"] = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[str]]: ... # docstr-coverage:excused `overload` @typing.overload @contextlib.contextmanager def open_inner_zipfile( zip_file: zipfile.ZipFile, inner_path: str | PurePath, *, operation: Operation = ..., representation: Literal["binary"] = ..., open_kwargs: Mapping[str, Any] | None = ..., encoding: str | None = ..., newline: str | None = ..., ) -> Generator[IO[bytes]]: ... @contextlib.contextmanager def open_inner_zipfile( zip_file: zipfile.ZipFile, inner_path: str | PurePath, *, operation: Operation = "read", representation: Representation = "text", open_kwargs: Mapping[str, Any] | None = None, encoding: str | None = None, newline: str | None = None, ) -> Generator[IO[str]] | Generator[IO[bytes]]: """Open a file inside an already opened zip archive.""" mode = _OPERATION_TO_UNQUALIFIED_MODE[operation] encoding = ensure_sensible_default_encoding(encoding, representation=representation) newline = ensure_sensible_newline(newline, representation=representation) # relative paths within zipfiles always follow POSIX, even on Windows inner_path = str(PurePosixPath(inner_path)) with ( zip_file.open(inner_path, mode=mode, **(open_kwargs or {})) as binary_file, _wrap_binary_if_needed( binary_file, representation, encoding=encoding, newline=newline ) as yf, ): yield yf ZZ = typing.TypeVar("ZZ", bound=IO[bytes]) @overload @contextlib.contextmanager def _wrap_binary_if_needed( file: ZZ, representation: Literal["text"], *, encoding: str | None = ..., newline: str | None = ..., ) -> Generator[io.TextIOWrapper[ZZ]]: ... @overload @contextlib.contextmanager def _wrap_binary_if_needed( file: ZZ, representation: Literal["binary"], *, encoding: str | None = ..., newline: str | None = ..., ) -> Generator[ZZ]: ... @contextlib.contextmanager def _wrap_binary_if_needed( file: ZZ, representation: Representation, *, encoding: str | None = None, newline: str | None = None, ) -> Generator[ZZ | io.TextIOWrapper[ZZ]]: if representation == "text": with io.TextIOWrapper(file, encoding=encoding, newline=newline) as text_file: yield text_file elif representation == "binary": yield file else: raise InvalidRepresentationError(representation) @contextlib.contextmanager def safe_open_dict_reader( f: TextSource, *, delimiter: str = "\t", **kwargs: Any ) -> Generator[csv.DictReader[str]]: """Open a CSV dictionary reader, wrapping :func:`csv.DictReader`. :param f: A path to a file, or an already open text-based IO object :param delimiter: The delimiter for writing to CSV :param kwargs: Keyword arguments to pass to :func:`csv.DictReader` :yields: A CSV reader object, constructed from :func:`csv.DictReader` """ with _open_read_text(f) as file: yield csv.DictReader(file, delimiter=delimiter, **kwargs) def is_url(obj: Any) -> TypeGuard[str]: """Check if the object is a URL.""" return isinstance(obj, str) and obj.startswith(("http://", "https://")) class UrlOpenKwargs(typing.TypedDict): """Keyword arguments for :func:`urllib.request.urlopen`.""" timeout: typing.NotRequired[int | None] # docstr-coverage:excused `overload` @overload @contextlib.contextmanager def open_url( url: str, *, representation: Literal["text"] = ..., **kwargs: Unpack[OpenKwargs], ) -> Generator[IO[str]]: ... # docstr-coverage:excused `overload` @overload @contextlib.contextmanager def open_url( url: str, *, representation: Literal["binary"] = ..., **kwargs: Unpack[OpenKwargs], ) -> Generator[IO[bytes]]: ... @contextlib.contextmanager def open_url( url: str, *, representation: Representation = "text", **kwargs: Unpack[OpenKwargs], ) -> Generator[IO[str]] | Generator[IO[bytes]]: """Get a file-like object from a URL.""" if (buffering := kwargs.get("buffering")) and buffering != -1: raise NotImplementedError("can not buffer when opening a URL") urlopen_kwargs: UrlOpenKwargs if timeout := kwargs.get("timeout"): urlopen_kwargs = {"timeout": timeout} else: urlopen_kwargs = {} with urllib.request.urlopen(url, **urlopen_kwargs) as response: # noqa:S310 match representation: case "text": yield io.TextIOWrapper( response, encoding=kwargs.get("encoding"), newline=kwargs.get("newline") ) case "binary": yield io.BufferedReader(response) case _: raise InvalidRepresentationError(representation) cthoyt-pystow-54a586e/src/pystow/utils/testing.py000066400000000000000000000007421525577235500223000ustar00rootroot00000000000000"""Utilities for testing.""" from collections.abc import Callable from typing import Any, TypeVar __all__ = [ "requires_package", ] _FT = TypeVar("_FT", bound=Callable[..., Any]) def requires_package(package: str) -> Callable[[_FT], _FT]: """Construct a decorator for a :class:`unittest.TestCase`.""" import importlib.util import unittest return unittest.skipUnless( importlib.util.find_spec(package), f"packages is not installed: {package}" ) cthoyt-pystow-54a586e/src/pystow/version.py000066400000000000000000000001271525577235500211450ustar00rootroot00000000000000"""Version information for PyStow.""" __all__ = [ "VERSION", ] VERSION = "0.9.8" cthoyt-pystow-54a586e/tests/000077500000000000000000000000001525577235500161145ustar00rootroot00000000000000cthoyt-pystow-54a586e/tests/__init__.py000066400000000000000000000000301525577235500202160ustar00rootroot00000000000000"""Tests for pystow.""" cthoyt-pystow-54a586e/tests/constants.py000066400000000000000000000006361525577235500205070ustar00rootroot00000000000000"""Test constants.""" from __future__ import annotations import os import unittest from pathlib import Path __all__ = [ "HERE", "RESOURCES", "skip_on_windows", ] skip_on_windows = unittest.skipIf( os.name == "nt", reason="Funny stuff happens in requests with a file adapter on windows that adds line breaks", ) HERE = Path(__file__).resolve().parent RESOURCES = HERE.joinpath("resources") cthoyt-pystow-54a586e/tests/resources/000077500000000000000000000000001525577235500201265ustar00rootroot00000000000000cthoyt-pystow-54a586e/tests/resources/test.html000066400000000000000000000000461525577235500217730ustar00rootroot00000000000000 test text cthoyt-pystow-54a586e/tests/resources/test.txt000066400000000000000000000000241525577235500216420ustar00rootroot00000000000000this is a test file cthoyt-pystow-54a586e/tests/resources/test.txt.bz2000066400000000000000000000000661525577235500223440ustar00rootroot00000000000000BZh91AY&SY6 р@#d " 42 N9HrE8P6 cthoyt-pystow-54a586e/tests/resources/test.txt.gz000066400000000000000000000000571525577235500222670ustar00rootroot00000000000000fRftest.txt+,VD̜T.Jcthoyt-pystow-54a586e/tests/resources/test.txt.md5000066400000000000000000000000401525577235500223240ustar00rootroot000000000000004221d002ceb5d3c9e9137e495ceaa647cthoyt-pystow-54a586e/tests/resources/test.txt.xz000066400000000000000000000001141525577235500223020ustar00rootroot000000000000007zXZִF!t/this is a test file ^tm, m}YZcthoyt-pystow-54a586e/tests/resources/test.txt.zst000066400000000000000000000000401525577235500224570ustar00rootroot00000000000000(/Xthis is a test file cthoyt-pystow-54a586e/tests/resources/test.xml000066400000000000000000000002171525577235500216270ustar00rootroot00000000000000 element 1 element 2 cthoyt-pystow-54a586e/tests/resources/test_1.csv000066400000000000000000000000521525577235500220370ustar00rootroot00000000000000h1,h2,h3 v1_1,v1_2,v1_3 v2_1,v2_2,v2_3 cthoyt-pystow-54a586e/tests/resources/test_1.json000066400000000000000000000000241525577235500222140ustar00rootroot00000000000000{ "key": "value" }cthoyt-pystow-54a586e/tests/resources/test_1.pkl000066400000000000000000000001171525577235500220340ustar00rootroot00000000000000D](h1h2h3v1_1v1_2v1_3v2_1v2_2v2_3e.cthoyt-pystow-54a586e/tests/resources/test_1.tsv000066400000000000000000000000471525577235500220640ustar00rootroot00000000000000h1 h2 h3 v1_1 v1_2 v1_3 v2_1 v2_2 v2_3 cthoyt-pystow-54a586e/tests/resources/test_verbose.txt.md5000066400000000000000000000000561525577235500240600ustar00rootroot00000000000000MD5(text.txt)=4221d002ceb5d3c9e9137e495ceaa647cthoyt-pystow-54a586e/tests/resources/test_wrong.txt.md5000066400000000000000000000000041525577235500235400ustar00rootroot00000000000000yolocthoyt-pystow-54a586e/tests/test_api.py000066400000000000000000000101641525577235500203000ustar00rootroot00000000000000"""Test for API completeness.""" from __future__ import annotations import inspect import unittest from collections.abc import Callable from typing import TypeVar, cast import pandas as pd import rdflib from lxml import etree import pystow from pystow import Module SKIP = {"__init__"} X = TypeVar("X") def _df_equal(a: pd.DataFrame, b: pd.DataFrame, msg: str | None = None) -> bool: return bool(a.values.tolist() == b.values.tolist()) def _rdf_equal(a: rdflib.Graph, b: rdflib.Graph, msg: str | None = None) -> bool: return {tuple(t) for t in a} == {tuple(t) for t in b} def _etree_equal(a: etree.ElementTree, b: etree.ElementTree, msg: str | None = None) -> bool: return cast(str, etree.tostring(a)) == cast(str, etree.tostring(b)) class TestExposed(unittest.TestCase): """Test API exposure.""" def setUp(self) -> None: """Set up the test case.""" self.addTypeEqualityFunc(pd.DataFrame, _df_equal) # type:ignore[arg-type] self.addTypeEqualityFunc(rdflib.Graph, _rdf_equal) # type:ignore[arg-type] self.addTypeEqualityFunc(type(etree.ElementTree(etree.Element("root"))), _etree_equal) # type:ignore[arg-type] def assert_io( self, obj: X, extension: str, dump: Callable[..., None], load: Callable[..., X] ) -> None: """Test an object can be dumped and loaded. :param obj: The object to dump :param extension: The extension to use :param dump: The dump function :param load: The load function """ name = f"test.{extension}" path = pystow.join("test", name=name) if path.is_file(): path.unlink() self.assertFalse(path.is_file()) dump("test", name=name, obj=obj) self.assertTrue(path.is_file()) self.assertEqual(obj, load("test", name=name)) def test_exposed(self) -> None: """Test that all module-level functions also have a counterpart in the top-level API.""" for name, func in Module.__dict__.items(): if not inspect.isfunction(func) or name in SKIP: continue with self.subTest(name=name): self.assertIn( name, pystow.api.__all__, msg=f"Module.{name} should be included in from `pystow.api.__all__`.", ) self.assertTrue( hasattr(pystow.api, name), msg=f"`Module.{name} should be exposed as a top-level " f"function in `pystow.api`.", ) self.assertTrue( hasattr(pystow, name), msg=f"`pystow.api.{name}` should be imported in `pystow.__init__`.", ) def test_io(self) -> None: """Test IO functions.""" obj = ["a", "b", "c"] for ext, dump, load in [ ("json", pystow.dump_json, pystow.load_json), ("pkl", pystow.dump_pickle, pystow.load_pickle), ]: with self.subTest(ext=ext): self.assert_io(obj, extension=ext, dump=dump, load=load) # type:ignore def test_pd_io(self) -> None: """Test pandas IO.""" columns = list("abc") data = [(1, 2, 3), (4, 5, 6)] df = pd.DataFrame(data, columns=columns) self.assert_io(df, extension="tsv", load=pystow.load_df, dump=pystow.dump_df) def test_rdf_io(self) -> None: """Test RDFlib IO.""" graph = rdflib.Graph() graph.add( ( rdflib.URIRef("http://example.com/subject"), rdflib.URIRef("http://example.com/predicate"), rdflib.URIRef("http://example.com/object"), ) ) self.assertEqual(1, len(graph)) self.assert_io(graph, extension="ttl", dump=pystow.dump_rdf, load=pystow.load_rdf) def test_xml_io(self) -> None: """Test XML I/O.""" root = etree.Element("root") root.set("interesting", "somewhat") etree.SubElement(root, "test") my_tree = etree.ElementTree(root) self.assert_io(my_tree, extension="xml", dump=pystow.dump_xml, load=pystow.load_xml) cthoyt-pystow-54a586e/tests/test_caching.py000066400000000000000000000121021525577235500211150ustar00rootroot00000000000000"""Tests for caching.""" from __future__ import annotations import os import tempfile import unittest from pathlib import Path from pydantic import BaseModel from pystow.cache import CachedPickle, CachedPydantic EXPECTED = 5 EXPECTED_2 = 6 class TestCache(unittest.TestCase): """Tests for caches.""" def setUp(self) -> None: """Set up the test case with a temporary directory.""" self.tmpdir = tempfile.TemporaryDirectory() self.directory = Path(self.tmpdir.name) def tearDown(self) -> None: """Tear down the test case's temporary directory.""" self.tmpdir.cleanup() def test_cache_exception(self) -> None: """Test that exceptions aren't swallowed.""" path = self.directory.joinpath("test.pkl") self.assertFalse(path.is_file()) @CachedPickle(path=path) def _f1() -> None: raise NotImplementedError self.assertFalse(path.is_file(), msg="function has not been called") with self.assertRaises(NotImplementedError): _f1() self.assertFalse( path.is_file(), msg="file should not have been created if an exception was thrown by the function", ) def test_cache_pickle(self) -> None: """Test caching a pickle.""" path = self.directory.joinpath("test.pkl") self.assertFalse( path.is_file(), msg="the file should not exist at the beginning of the test", ) raise_flag = True @CachedPickle(path=path) def _f1() -> int: if raise_flag: raise ValueError return EXPECTED self.assertFalse(path.is_file(), msg="the file should not exist until function is called") with self.assertRaises(ValueError): _f1() self.assertFalse( path.is_file(), msg="the function should throw an exception " "because of the flag, and no file should be created", ) raise_flag = False actual = _f1() self.assertEqual(EXPECTED, actual) self.assertTrue(path.is_file(), msg="a file should have been created") raise_flag = True actual_2 = _f1() # if raises, the caching mechanism didn't work self.assertEqual(EXPECTED, actual_2) self.assertTrue(path.is_file()) os.unlink(path) self.assertFalse(path.is_file()) with self.assertRaises(ValueError): _f1() @CachedPickle(path=path, force=True) def _f2() -> int: return EXPECTED_2 self.assertEqual(EXPECTED_2, _f2()) # overwrites the file self.assertEqual(EXPECTED_2, _f1()) def test_cache_pydantic(self) -> None: """Test caching a pydantic.""" path = self.directory.joinpath("test.json") self.assertFalse( path.is_file(), msg="the file should not exist at the beginning of the test", ) class Model(BaseModel): """A simple model.""" value: int raise_flag = True @CachedPydantic(path=path, model_cls=Model) def _f1() -> Model: if raise_flag: raise ValueError return Model(value=EXPECTED) self.assertFalse(path.is_file(), msg="the file should not exist until function is called") with self.assertRaises(ValueError): _f1() self.assertFalse( path.is_file(), msg="the function should throw an exception " "because of the flag, and no file should be created", ) raise_flag = False actual = _f1() self.assertIsInstance(actual, Model) self.assertEqual(EXPECTED, actual.value) self.assertTrue(path.is_file(), msg="a file should have been created") raise_flag = True actual_2 = _f1() # if raises, the caching mechanism didn't work self.assertIsInstance(actual_2, Model) self.assertEqual(EXPECTED, actual_2.value) self.assertTrue(path.is_file()) os.unlink(path) self.assertFalse(path.is_file()) with self.assertRaises(ValueError): _f1() @CachedPydantic(path=path, model_cls=Model, force=True) def _f2() -> Model: return Model(value=EXPECTED_2) self.assertEqual(EXPECTED_2, _f2().value) # overwrites the file self.assertEqual(EXPECTED_2, _f1().value) def test_no_cache(self) -> None: """Test that no caching happens.""" path = self.directory.joinpath("test.pkl") sentinel_value = 5 self.assertFalse(path.is_file()) @CachedPickle(path=path, cache=False) def _f1() -> int: return sentinel_value self.assertFalse(path.is_file(), msg="function has not been called") # check the following twice, just for good measure! for _ in range(2): self.assertEqual(sentinel_value, _f1()) self.assertFalse( path.is_file(), msg="file should not have been created since caching was turned off", ) cthoyt-pystow-54a586e/tests/test_config.py000066400000000000000000000105111525577235500207700ustar00rootroot00000000000000"""Test configuration loading.""" from __future__ import annotations import tempfile import unittest from configparser import ConfigParser from pathlib import Path from typing import ClassVar import pystow from pystow.config_api import CONFIG_HOME_ENVVAR, _get_cfp from pystow.utils import mock_envvar class TestConfig(unittest.TestCase): """Test configuration.""" test_section: ClassVar[str] test_option: ClassVar[str] test_value: ClassVar[str] cfp: ClassVar[ConfigParser] @classmethod def setUpClass(cls) -> None: """Set up the class for testing.""" cls.test_section = "test" cls.test_option = "option" cls.test_value = "value" cls.cfp = _get_cfp(cls.test_section) cls.cfp.add_section(cls.test_section) cls.cfp.set( section=cls.test_section, option=cls.test_option, value=cls.test_value, ) def test_env_cast(self) -> None: """Test casting works properly when getting from the environment.""" with mock_envvar("TEST_VAR", "1234"): self.assertEqual("1234", pystow.get_config("test", "var")) self.assertEqual("1234", pystow.get_config("test", "var", dtype=str)) self.assertEqual(1234, pystow.get_config("test", "var", dtype=int)) with self.assertRaises(ValueError): pystow.get_config("test", "var", dtype=bool) with self.assertRaises(TypeError): pystow.get_config("test", "var", dtype=object) def test_get_config(self) -> None: """Test lookup not existing.""" self.assertIsNone(pystow.get_config(self.test_section, "key")) self.assertEqual("1234", pystow.get_config(self.test_section, "key", default="1234")) value = "not_value" self.assertEqual( value, pystow.get_config(self.test_section, self.test_option, passthrough=value) ) self.assertEqual(1, pystow.get_config(self.test_section, self.test_option, passthrough=1)) self.assertEqual( 1, pystow.get_config(self.test_section, self.test_option, passthrough="1", dtype=int) ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough="1", dtype=bool), ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough="yes", dtype=bool), ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough="Yes", dtype=bool), ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough="YES", dtype=bool), ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough="True", dtype=bool), ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough="TRUE", dtype=bool), ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough="T", dtype=bool), ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough="t", dtype=bool), ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough=True, dtype=bool), ) self.assertEqual( True, pystow.get_config(self.test_section, self.test_option, passthrough=1, dtype=bool) ) def test_subsection(self) -> None: """Test subsections.""" with tempfile.TemporaryDirectory() as directory, mock_envvar(CONFIG_HOME_ENVVAR, directory): directory_ = Path(directory) path = directory_.joinpath("test.ini") self.assertFalse(path.is_file(), msg="file should not already exist") self.assertIsNone(pystow.get_config("test:subtest", "key")) self.assertFalse(path.is_file(), msg="getting config should not create a file") pystow.write_config("test:subtest", "key", "value") self.assertTrue(path.is_file(), msg=f"{list(directory_.iterdir())}") self.assertEqual("value", pystow.get_config("test:subtest", "key")) cthoyt-pystow-54a586e/tests/test_git_utils.py000066400000000000000000000023331525577235500215310ustar00rootroot00000000000000"""Test git utilities.""" import unittest from pystow.git import ( clone_github_tempdir, create_branch, get_current_branch, has_local_branch, is_likely_default_branch, ) from tests.constants import skip_on_windows @skip_on_windows class TestGitUtils(unittest.TestCase): """Test git utilities.""" def test_temporary_github_clone(self) -> None: """Test temporarily cloning a GitHub repository.""" name = "test-branch-name" with clone_github_tempdir("cthoyt", "cthoyt") as directory: self.assertTrue(directory.joinpath("README.md").is_file()) self.assertEqual("master", get_current_branch(directory)) self.assertTrue(is_likely_default_branch(directory)) self.assertFalse(has_local_branch(directory, name)) create_branch(directory, name) self.assertEqual(name, get_current_branch(directory)) self.assertTrue(has_local_branch(directory, name)) self.assertFalse(is_likely_default_branch(directory)) # Test if we do this again, the branch didn't persist with clone_github_tempdir("cthoyt", "cthoyt") as directory: self.assertFalse(has_local_branch(directory, name)) cthoyt-pystow-54a586e/tests/test_graph.py000066400000000000000000000037651525577235500206410ustar00rootroot00000000000000"""Test graph.""" import tempfile import unittest from pystow.graph import GraphCachePaths, build_graph_cache from tests.constants import skip_on_windows @skip_on_windows class TestGraph(unittest.TestCase): """Test graph.""" def test_cache(self) -> None: """Test error on missing directory.""" with self.assertRaises(NotADirectoryError): GraphCachePaths.from_directory("blahblahblah") def test_graph(self) -> None: """Test building a graph.""" edges = [ ("a", "b"), ("a", "c"), ("a", "d"), ("b", "d"), ("d", "e"), ] with tempfile.TemporaryDirectory() as tmpdir: with self.assertRaises(TypeError): build_graph_cache(edges, tmpdir) # type:ignore[arg-type] paths = GraphCachePaths.from_directory(tmpdir) self.assertFalse(paths.exists()) graph = build_graph_cache(lambda: edges, paths, sort_nodes=True, progress=False) self.assertTrue(paths.exists()) self.assertEqual([], graph.out_edges("ZZZ")) self.assertEqual([], graph.in_edges("ZZZ")) with self.assertRaises(KeyError): graph.out_edges("ZZZ", raise_on_missing=True) with self.assertRaises(KeyError): graph.in_edges("ZZZ", raise_on_missing=True) self.assertEqual({"b", "c", "d"}, set(graph.out_edges("a"))) self.assertEqual({"d"}, set(graph.out_edges("b"))) self.assertEqual(set(), set(graph.out_edges("c"))) self.assertEqual({"e"}, set(graph.out_edges("d"))) self.assertEqual(set(), set(graph.out_edges("e"))) self.assertEqual(set(), set(graph.in_edges("a"))) self.assertEqual({"a"}, set(graph.in_edges("b"))) self.assertEqual({"a"}, set(graph.in_edges("c"))) self.assertEqual({"a", "b"}, set(graph.in_edges("d"))) self.assertEqual({"d"}, set(graph.in_edges("e"))) cthoyt-pystow-54a586e/tests/test_module.py000066400000000000000000000333411525577235500210160ustar00rootroot00000000000000"""Tests for PyStow.""" from __future__ import annotations import bz2 import contextlib import itertools as itt import json import lzma import os import pickle import shutil import tempfile import unittest from collections.abc import Generator, Mapping from pathlib import Path from typing import Any from unittest import mock import pandas as pd import yaml import pystow from pystow import join from pystow.constants import PYSTOW_HOME_ENVVAR, PYSTOW_NAME_ENVVAR from pystow.impl import Module from pystow.utils import ( get_home, get_name, mock_envvar, n, write_pickle_gz, write_sql, write_tarfile_csv, write_zipfile_csv, ) HERE = Path(__file__).parent.resolve() RESOURCES = HERE.joinpath("resources") TSV_NAME = "test_1.tsv" TSV_URL = f"{n()}/{TSV_NAME}" HTML_NAME = "test.html" HTML_URL = f"{n()}/{HTML_NAME}" HTML_PATH = RESOURCES / HTML_NAME SQLITE_NAME = "test_1.db" SQLITE_URL = f"{n()}/{SQLITE_NAME}" SQLITE_PATH = RESOURCES / SQLITE_NAME SQLITE_TABLE = "testtable" JSON_NAME = "test_1.json" JSON_URL = f"{n()}/{JSON_NAME}" JSON_PATH = RESOURCES / JSON_NAME YAML_NAME = "test_1.yaml" YAML_URL = f"{n()}/{YAML_NAME}" YAML_PATH = RESOURCES / YAML_NAME PICKLE_NAME = "test_1.pkl" PICKLE_URL = f"{n()}/{PICKLE_NAME}" PICKLE_PATH = RESOURCES / PICKLE_NAME PICKLE_GZ_NAME = "test_1.pkl.gz" PICKLE_GZ_URL = f"{n()}/{PICKLE_GZ_NAME}" PICKLE_GZ_PATH = RESOURCES / PICKLE_GZ_NAME JSON_BZ2_NAME = "test_1.json.bz2" JSON_BZ2_URL = f"{n()}/{JSON_BZ2_NAME}" JSON_BZ2_PATH = RESOURCES / JSON_BZ2_NAME MOCK_FILES: Mapping[str, Path] = { TSV_URL: RESOURCES / TSV_NAME, JSON_URL: JSON_PATH, YAML_URL: YAML_PATH, JSON_BZ2_URL: JSON_BZ2_PATH, PICKLE_URL: PICKLE_PATH, PICKLE_GZ_URL: PICKLE_GZ_PATH, SQLITE_URL: SQLITE_PATH, HTML_URL: HTML_PATH, } TEST_TSV_ROWS = [ ("h1", "h2", "h3"), ("v1_1", "v1_2", "v1_3"), ("v2_1", "v2_2", "v2_3"), ] TEST_DF = pd.DataFrame(TEST_TSV_ROWS) TEST_JSON = TEST_YAML = {"key": "value"} # Make the pickle file if not PICKLE_PATH.is_file(): PICKLE_PATH.write_bytes(pickle.dumps(TEST_TSV_ROWS)) if not SQLITE_PATH.is_file(): write_sql(TEST_DF, name=SQLITE_TABLE, path=SQLITE_PATH, index=False) if not JSON_PATH.is_file(): JSON_PATH.write_text(json.dumps(TEST_JSON)) if not YAML_PATH.is_file(): YAML_PATH.write_text(yaml.safe_dump(TEST_YAML)) if not JSON_BZ2_PATH.is_file(): with bz2.open(JSON_BZ2_PATH, mode="wt") as file: json.dump(TEST_JSON, file, indent=2) class TestMocks(unittest.TestCase): """Tests for :mod:`pystow` mocks and context managers.""" def test_mock_home(self) -> None: """Test that home can be properly mocked.""" name = n() with tempfile.TemporaryDirectory() as d: expected_path = Path(d) / name self.assertFalse(expected_path.exists()) with mock_envvar(PYSTOW_HOME_ENVVAR, expected_path.as_posix()): self.assertFalse(expected_path.exists()) self.assertEqual(expected_path, get_home(ensure_exists=False)) self.assertFalse(expected_path.exists()) def test_mock_name(self) -> None: """Test that the name can be properly mocked.""" name = n() expected_path = Path.home() / name self.assertFalse(expected_path.exists()) with mock_envvar(PYSTOW_NAME_ENVVAR, name): self.assertEqual(name, get_name()) self.assertFalse(expected_path.exists()) self.assertEqual(expected_path, get_home(ensure_exists=False)) self.assertFalse(expected_path.exists()) class TestJoin(unittest.TestCase): """Tests for :mod:`pystow`.""" def setUp(self) -> None: """Set up the test case.""" self.directory = tempfile.TemporaryDirectory() def tearDown(self) -> None: """Tear down the test case.""" self.directory.cleanup() @contextlib.contextmanager def mock_directory(self) -> Generator[Path, None, None]: """Use this test case's temporary directory as a mock environment variable. :yield: The mock directory's path """ with mock_envvar(PYSTOW_HOME_ENVVAR, self.directory.name): yield Path(self.directory.name) @staticmethod def mock_download() -> Any: """Mock connection to the internet using local resource files. :returns: A patch object that can be applied to the pystow download function """ def _mock_get_data(url: str, path: str | Path, **_kwargs: Any) -> Path: return Path(shutil.copy(MOCK_FILES[url], path)) return mock.patch("pystow.utils.download", side_effect=_mock_get_data) @staticmethod def mock_download_once(local_path: str | Path) -> Any: """Mock connection to the internet using local resource files. :param local_path: the path to the file to mock :returns: A patch object that can be applied to the pystow download function """ def _mock_get_data(path: str | Path, **_kwargs: Any) -> Path: return Path(shutil.copy(local_path, path)) return mock.patch("pystow.utils.download", side_effect=_mock_get_data) def join(self, *parts: str) -> Path: """Help join the parts to this test case's temporary directory. :param parts: The file path parts that are joined with this test case's directory :returns: A path to the file """ return Path(self.directory.name).joinpath(*parts) def test_mock(self) -> None: """Test that mocking the directory works properly for this test case.""" with self.mock_directory(): self.assertEqual(os.getenv(PYSTOW_HOME_ENVVAR), self.directory.name) def test_join(self) -> None: """Test the :func:`pystow.join` function.""" parts_examples = [ [n()], [n(), n()], [n(), n(), n()], ] with self.mock_directory(): for parts in parts_examples: with self.subTest(parts=parts): self.assertEqual(self.join(*parts), join(*parts)) def test_join_with_version(self) -> None: """Test the join function when a version is present.""" with self.mock_directory(): key = "key" version = "v1" self.assertEqual( self.join(key, version), pystow.join(key, version=version), ) parts = [n()] self.assertEqual( self.join(key, version, *parts), pystow.join(key, *parts, version=version) ) parts = [n()] name = "yup.tsv" self.assertEqual( self.join(key, version, *parts, name), pystow.join(key, *parts, version=version, name=name), ) def _version_getter() -> str: return "v2" parts = [n()] name = "yup.tsv" self.assertEqual( self.join(key, _version_getter(), *parts, name), pystow.join(key, *parts, version=_version_getter, name=name), ) with self.assertRaises(ValueError): pystow.join(key, version="/") def test_ensure(self) -> None: """Test ensuring various files.""" write_pickle_gz(TEST_TSV_ROWS, path=PICKLE_GZ_PATH) with self.mock_directory(), self.mock_download(): with self.subTest(type="tsv"): df = pystow.ensure_csv("test", url=TSV_URL, download_kwargs={"backend": "urllib"}) self.assertEqual(3, len(df.columns)) df2 = pystow.load_df("test", name=TSV_NAME) self.assertEqual(df.values.tolist(), df2.values.tolist()) with self.subTest(type="json"): j = pystow.ensure_json("test", url=JSON_URL) self.assertEqual(TEST_JSON, j) j2 = pystow.load_json("test", name=JSON_NAME) self.assertEqual(j, j2) with self.subTest(type="yaml"): j = pystow.ensure_yaml("test", url=YAML_URL) self.assertEqual(TEST_YAML, j) j2 = pystow.load_yaml("test", name=YAML_NAME) self.assertEqual(j, j2) with self.subTest(type="pickle"): p = pystow.ensure_pickle("test", url=PICKLE_URL) self.assertEqual(3, len(p)) p2 = pystow.load_pickle("test", name=PICKLE_NAME) self.assertEqual(p, p2) with self.subTest(type="pickle_gz"): p = pystow.ensure_pickle_gz("test", url=PICKLE_GZ_URL) self.assertEqual(3, len(p)) p2 = pystow.load_pickle_gz("test", name=PICKLE_GZ_NAME) self.assertEqual(p, p2) with self.subTest(type="json_bz2"): p = pystow.ensure_json_bz2("test", url=JSON_BZ2_URL) self.assertEqual(TEST_JSON, p) with self.subTest(type="html"): soup = pystow.ensure_soup("test", url=HTML_URL) self.assertEqual("test text", soup.text.strip()) def test_open_fail(self) -> None: """Test opening a missing file.""" with self.assertRaises(FileNotFoundError), pystow.open("nope", name="nope"): pass with self.assertRaises(FileNotFoundError): pystow.load_json("nope", name="nope") def test_ensure_open_lzma(self) -> None: """Test opening lzma-encoded files.""" with tempfile.TemporaryDirectory() as directory, self.mock_directory(): path = Path(directory) / n() with self.mock_download_once(path): with lzma.open(path, "wt") as file_1: for row in TEST_TSV_ROWS: print(*row, sep="\t", file=file_1) # FIXME this ignore needs to be removed and addressed with pystow.ensure_open_lzma("test", url=n()) as file_2: # type: ignore df = pd.read_csv(file_2, sep="\t") self.assertEqual(3, len(df.columns)) def test_ensure_open_zip(self) -> None: """Test opening tar-encoded files.""" with tempfile.TemporaryDirectory() as directory, self.mock_directory(): path = Path(directory) / n() inner_path = n() with self.mock_download_once(path): write_zipfile_csv(TEST_DF, path, inner_path) # test many modes, luckily, because pandas can handle everything with pystow.ensure_open_zip( "test", url=n(), inner_path=inner_path, mode="rb" ) as file: df = pd.read_csv(file, sep="\t") self.assertEqual(3, len(df.columns)) with pystow.ensure_open_zip( "test", url=n(), inner_path=inner_path, mode="rt" ) as file_text: df = pd.read_csv(file_text, sep="\t") self.assertEqual(3, len(df.columns)) def test_ensure_open_tarfile(self) -> None: """Test opening tarfile-encoded files.""" with tempfile.TemporaryDirectory() as directory, self.mock_directory(): path = Path(directory) / n() inner_path = n() with self.mock_download_once(path): write_tarfile_csv(TEST_DF, path, inner_path) with pystow.ensure_open_tarfile("test", url=n(), inner_path=inner_path) as file: df = pd.read_csv(file, sep="\t") self.assertEqual(3, len(df.columns)) def test_ensure_module(self) -> None: """Test that the ``ensure_exist`` argument in :meth:`Module.from_key` works properly.""" parts_examples = [ [n()], [n(), n()], [n(), n(), n()], ] ensure_examples = [False, True] for ensure_exists, parts in itt.product(ensure_examples, parts_examples): with self.subTest(ensure_exists=ensure_exists, parts=parts), self.mock_directory(): expected_directory = self.join(*parts) module = Module.from_key(*parts, ensure_exists=ensure_exists) self.assertEqual(expected_directory, module.base) self.assertIs( expected_directory.exists(), ensure_exists, msg=f"{expected_directory} should{'' if ensure_exists else ' not'} exist.", ) def test_ensure_custom(self) -> None: """Test ensure with custom provider.""" with self.mock_directory(): # create a minimal provider def touch_file(path: Path, **_kwargs: Any) -> None: """Create a file. :param path: the file path :param _kwargs: ignored keywords """ path.touch() # wrap to record calls provider = mock.Mock(wraps=touch_file) # the keyword-based parameters for the provider kwargs: dict[str, Any] = {"a": 4, "c": {0: 1, 5: 7}} # call first time name = n() path = pystow.ensure_custom("test", name=name, provider=provider, **kwargs) self.assertTrue(path.is_file()) # call a second time path = pystow.ensure_custom("test", name=name, provider=provider, **kwargs) # ensure that the provider was only called once with the given parameters provider.assert_called_once_with(path, **kwargs) def test_ensure_open_sqlite(self) -> None: """Test caching SQLite.""" with ( self.mock_directory(), self.mock_download(), pystow.ensure_open_sqlite("test", url=SQLITE_URL) as conn, ): df = pd.read_sql(f"SELECT * from {SQLITE_TABLE}", conn) # noqa:S608 self.assertEqual(3, len(df.columns)) cthoyt-pystow-54a586e/tests/test_utils/000077500000000000000000000000001525577235500203135ustar00rootroot00000000000000cthoyt-pystow-54a586e/tests/test_utils/__init__.py000066400000000000000000000000331525577235500224200ustar00rootroot00000000000000"""Tests for utilities.""" cthoyt-pystow-54a586e/tests/test_utils/test_pydantic_utils.py000066400000000000000000000127171525577235500247670ustar00rootroot00000000000000"""Tests for Pydantic utilities.""" import datetime import json import tempfile import unittest from pathlib import Path from typing import Any from pystow.utils import ( read_pydantic_jsonl, read_pydantic_tsv, requires_package, safe_open_writer, stream_write_pydantic_jsonl, write_pydantic_json, write_pydantic_jsonl, write_pydantic_yaml, ) @requires_package("pydantic") class TestPydanticUtils(unittest.TestCase): """Tests for Pydantic utilities.""" def setUp(self) -> None: """Set up the test case.""" self._temporary_directory = tempfile.TemporaryDirectory() self.directory = Path(self._temporary_directory.name) def tearDown(self) -> None: """Clean up the test case.""" self._temporary_directory.cleanup() def test_pydantic_jsonl(self) -> None: """Test writing Pydantic.""" from pydantic import BaseModel class Model(BaseModel): """A test model.""" name: str models = [Model(name=f"test {i}") for i in range(3)] path = self.directory.joinpath("data.jsonl") write_pydantic_jsonl(models, path) self.assertEqual(models, read_pydantic_jsonl(path, Model)) def test_streaming_jsonl_writer(self) -> None: """Test writing Pydantic.""" from pydantic import BaseModel class Model(BaseModel): """A test model.""" name: str models = [Model(name=f"test {i}") for i in range(3)] path = self.directory.joinpath("data.jsonl") new_models = list(stream_write_pydantic_jsonl(models, path)) self.assertEqual(models, new_models) self.assertEqual(models, read_pydantic_jsonl(path, Model)) def test_read_jsonl_error_action(self) -> None: """Test error actions.""" import pydantic from pydantic import BaseModel class Model(BaseModel): """A test model.""" name: str path = self.directory.joinpath("data.jsonl") with path.open("w") as file: print(json.dumps({"name": "1"}), file=file) print(json.dumps({"name": "2"}), file=file) print(json.dumps({"nope": "3"}), file=file) self.assertEqual( [Model(name="1"), Model(name="2")], list(read_pydantic_jsonl(path, Model, failure_action="skip")), ) with self.assertRaises(pydantic.ValidationError): (list(read_pydantic_jsonl(path, Model, failure_action="raise")),) def test_read_tsv(self) -> None: """Test reading pydantic models from a TSV file.""" from pydantic import BaseModel class Model(BaseModel): """A test model.""" name: str date: datetime.date path = self.directory.joinpath("data.tsv") with safe_open_writer(path) as writer: writer.writerow(("name", "date")) writer.writerow(("1", "2025-01-01")) writer.writerow(("2", "2025-01-02")) writer.writerow(("3", "2025-01-03")) self.assertEqual( [ Model(name="1", date=datetime.date.fromisoformat("2025-01-01")), Model(name="2", date=datetime.date.fromisoformat("2025-01-02")), Model(name="3", date=datetime.date.fromisoformat("2025-01-03")), ], read_pydantic_tsv(path, Model), ) def test_read_tsv_with_processing(self) -> None: """Test reading pydantic models from a TSV file with custom row-based processing.""" from pydantic import BaseModel class Model(BaseModel): """A test model.""" name: str date: datetime.date def _process(record: dict[str, Any]) -> dict[str, Any]: date_value = record.pop("date") if len(date_value) == 4: record["date"] = f"{date_value}-01-01" else: record["date"] = date_value return record path = self.directory.joinpath("data.tsv") with safe_open_writer(path) as writer: writer.writerow(("name", "date")) writer.writerow(("1", "2025")) writer.writerow(("2", "2025-01-02")) writer.writerow(("3", "2025-01-03")) self.assertEqual( [ Model(name="1", date=datetime.date.fromisoformat("2025-01-01")), Model(name="2", date=datetime.date.fromisoformat("2025-01-02")), Model(name="3", date=datetime.date.fromisoformat("2025-01-03")), ], read_pydantic_tsv(path, Model, process=_process), ) def test_write_json(self) -> None: """Test writing pydantic models to a JSON file.""" from pydantic import BaseModel class Model(BaseModel): """A model with a name.""" name: str path = self.directory.joinpath("data.json") model = Model(name="test") write_pydantic_json(model, path) self.assertEqual("""{"name": "test"}\n""", path.read_text()) def test_write_yaml(self) -> None: """Test writing pydantic models to a YAML file.""" from pydantic import BaseModel class Model(BaseModel): """A model with a name.""" name: str path = self.directory.joinpath("data.json") model = Model(name="test") write_pydantic_yaml(model, path) self.assertEqual( "name: test\n", path.read_text(), ) cthoyt-pystow-54a586e/tests/test_utils/test_reiter.py000066400000000000000000000012751525577235500232230ustar00rootroot00000000000000"""Tests for itertools.""" import unittest from collections.abc import Iterable from pystow.utils import reyield class TestItertools(unittest.TestCase): """Tests for itertools.""" def test_reyield(self) -> None: """Test making a function re-iterable.""" sum_value = 0 square_sum_value = 0 def _accumulate(inner_elements: Iterable[int]) -> None: nonlocal sum_value for inner_element in inner_elements: sum_value += inner_element for element in reyield(_accumulate, range(5)): square_sum_value += element**2 self.assertEqual(10, sum_value) self.assertEqual(30, square_sum_value) cthoyt-pystow-54a586e/tests/test_utils/test_utils.py000066400000000000000000001045531525577235500230740ustar00rootroot00000000000000"""Tests for utilities.""" from __future__ import annotations import hashlib import itertools as itt import os import tarfile import tempfile import unittest import zipfile from pathlib import Path import numpy as np import pandas as pd import rdflib import requests from lxml import etree from requests_file import FileAdapter from pystow import utils from pystow.github import get_default_branch, get_repository from pystow.utils import ( DownloadError, HexDigestError, download, get_hash_hexdigest, get_hexdigests_remote, getenv_path, gunzip, gzip_compress, is_url, iter_tarred_csvs, iter_zipped_csvs, mkdir, mock_envvar, n, name_from_url, open_tarfile, open_url, open_zip_reader, open_zip_writer, open_zipfile, read_lzma_csv, read_tarfile_csv, read_tarfile_xml, read_zip_np, read_zipfile_csv, read_zipfile_rdf, read_zipfile_xml, safe_open, safe_open_dict_reader, safe_open_json, safe_open_reader, safe_open_writer, safe_open_yaml, safe_read_text, safe_write_text, tarfile_writestr, write_lzma_csv, write_tarfile_csv, write_tarfile_xml, write_zipfile_csv, write_zipfile_np, write_zipfile_rdf, write_zipfile_xml, ) from pystow.utils.safe_open import COMPRESSION_EXTENSIONS from tests.constants import RESOURCES TEST_TXT = RESOURCES.joinpath("test.txt") TEST_TXT_CONTENT = "this is a test file\n" TEST_TXT_MD5 = RESOURCES.joinpath("test.txt.md5") TEST_TXT_GZ = RESOURCES.joinpath("test.txt.gz") TEST_TXT_BZ2 = RESOURCES.joinpath("test.txt.bz2") TEST_TXT_LZMA = RESOURCES.joinpath("test.txt.xz") TEST_TXT_ZSTD = RESOURCES.joinpath("test.txt.zst") TEST_TXT_VERBOSE_MD5 = RESOURCES.joinpath("test_verbose.txt.md5") TEST_TXT_WRONG_MD5 = RESOURCES.joinpath("test_wrong.txt.md5") TEST_XML = RESOURCES.joinpath("test.xml") TEST_PATHS = [TEST_TXT, TEST_TXT_GZ, TEST_TXT_BZ2, TEST_TXT_LZMA, TEST_TXT_ZSTD] class _Session(requests.sessions.Session): """A mock session.""" def __init__(self) -> None: """Instantiate the patched session with an additional file adapter.""" super().__init__() self.mount("file://", FileAdapter()) requests.sessions.Session = _Session # type: ignore class TestUtils(unittest.TestCase): """Test utility functions.""" def test_name_from_url(self) -> None: """Test :func:`name_from_url`.""" data = [ ("test.tsv", "https://example.com/test.tsv"), ("test.tsv", "https://example.com/deeper/test.tsv"), ("test.tsv.gz", "https://example.com/deeper/test.tsv.gz"), ] for name, url in data: with self.subTest(name=name, url=url): self.assertEqual(name, name_from_url(url)) def test_file_values(self) -> None: """Test encodings.""" for url, value in [ (TEST_TXT, TEST_TXT_CONTENT), # `this is a test file\n` (TEST_TXT_MD5, "4221d002ceb5d3c9e9137e495ceaa647"), (TEST_TXT_VERBOSE_MD5, "MD5(text.txt)=4221d002ceb5d3c9e9137e495ceaa647"), (TEST_TXT_WRONG_MD5, "yolo"), ]: with self.subTest(name=url.name): self.assertEqual(value, requests.get(url.as_uri(), timeout=15).text) def test_get_hash(self) -> None: """Test directly calculating a hash digest.""" self.assertEqual( TEST_TXT_MD5.read_text(encoding="utf-8"), get_hash_hexdigest(TEST_TXT, "md5") ) def test_mkdir(self) -> None: """Test for ensuring a directory.""" with tempfile.TemporaryDirectory() as directory: directory_ = Path(directory) subdirectory = directory_ / "sd1" self.assertFalse(subdirectory.exists()) mkdir(subdirectory, ensure_exists=False) self.assertFalse(subdirectory.exists()) mkdir(subdirectory, ensure_exists=True) self.assertTrue(subdirectory.exists()) def test_mock_envvar(self) -> None: """Test that environment variables can be mocked properly.""" name, value = n(), n() self.assertNotIn(name, os.environ) with mock_envvar(name, value): self.assertIn(name, os.environ) self.assertEqual(value, os.getenv(name)) self.assertNotIn(name, os.environ) def test_getenv_path(self) -> None: """Test that :func:`getenv_path` works properly.""" envvar = n() with tempfile.TemporaryDirectory() as directory: directory_ = Path(directory) value = directory_ / n() default = directory_ / n() self.assertEqual(default, getenv_path(envvar, default)) with mock_envvar(envvar, value.as_posix()): self.assertEqual(value, getenv_path(envvar, default)) # Check that it goes back self.assertEqual(default, getenv_path(envvar, default)) def test_compressed_io(self) -> None: """Test that the read/write to compressed folder functions work.""" rows = [[1, 2], [3, 4], [5, 6]] columns = ["A", "B"] df = pd.DataFrame(rows, columns=columns) inner_path = "okay.tsv" data = [ ("test.zip", write_zipfile_csv, read_zipfile_csv), ("test.tar.gz", write_tarfile_csv, read_tarfile_csv), ] for name, writer, reader in data: with self.subTest(name=name), tempfile.TemporaryDirectory() as directory: path = Path(directory) / name self.assertFalse(path.exists()) writer(df, path=path, inner_path=inner_path) self.assertTrue(path.exists()) new_df = reader(path=path, inner_path=inner_path) self.assertEqual(list(df.columns), list(new_df.columns)) self.assertEqual(df.values.tolist(), new_df.values.tolist()) def test_xml_io(self) -> None: """Test that read/write for XML element tree works.""" root = etree.Element("Doc") level1 = etree.SubElement(root, "S") main = etree.SubElement(level1, "Text") main.text = "Thanks for contributing an answer to Stack Overflow!" second = etree.SubElement(level1, "Tokens") level2 = etree.SubElement(second, "Token", word="low") level3 = etree.SubElement(level2, "Morph") second1 = etree.SubElement(level3, "Lemma") second1.text = "sdfs" second1 = etree.SubElement(level3, "info") second1.text = "qw" level4 = etree.SubElement(level3, "Aff") second1 = etree.SubElement(level4, "Type") second1.text = "sdfs" second1 = etree.SubElement(level4, "Suf") second1.text = "qw" tree = etree.ElementTree(root) inner_path = "okay.tsv" data = [ ("test.zip", write_zipfile_xml, read_zipfile_xml), ("test.tar.gz", write_tarfile_xml, read_tarfile_xml), ] for name, writer, reader in data: with self.subTest(name=name), tempfile.TemporaryDirectory() as directory: path = Path(directory) / name self.assertFalse(path.exists()) writer(tree, path=path, inner_path=inner_path) self.assertTrue(path.exists()) new_tree = reader(path=path, inner_path=inner_path) self.assertEqual( etree.tostring(tree, pretty_print=True), etree.tostring(new_tree, pretty_print=True), ) def test_read_xml(self) -> None: """Test reading XML.""" root = utils.get_xml_root(TEST_XML) self.assertEqual(2, len(root), msg=f"elements: {list(root)}") elements = list(utils.iterparse_xml(TEST_XML, tag="element")) self.assertEqual(2, len(elements), msg=f"elements: {elements}") def test_numpy_io(self) -> None: """Test IO with numpy.""" arr = np.array([[0, 1], [2, 3]]) inner_path = "okay.npz" with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "test.zip" write_zipfile_np(arr, inner_path=inner_path, path=path) reloaded_arr = read_zip_np(path=path, inner_path=inner_path) self.assertTrue(np.array_equal(arr, reloaded_arr)) def test_rdf(self) -> None: """Test reading and writing RDF.""" t = ( rdflib.URIRef("https://orcid.org/0000-0003-4423-4370"), rdflib.RDF.type, rdflib.URIRef("http://purl.obolibrary.org/obo/NCBITaxon_9606"), ) graph = rdflib.Graph() graph.add(t) def _test_equal(gg: rdflib.Graph) -> None: records = list(gg.query("SELECT ?s ?p ?o WHERE {?s ?p ?o}")) self.assertEqual(1, len(records)) self.assertEqual(t, records[0]) with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "test.zip" inner_path = "test_graph.ttl" write_zipfile_rdf(graph, path, inner_path) g2 = read_zipfile_rdf(path, inner_path) _test_equal(g2) for name in ["test.ttl", "test.ttl.gz"]: path = Path(directory) / name utils.write_rdflib(graph, path) g3 = utils.read_rdflib(path) _test_equal(g3) # test text mode works path = Path(directory) / "test-2.ttl" with path.open("wt") as file: utils.write_rdflib(graph, file) _test_equal(utils.read_rdflib(path)) # test binary mode works path = Path(directory) / "test-3.ttl" with path.open("wb") as bfile: utils.write_rdflib(graph, bfile) _test_equal(utils.read_rdflib(path)) def test_safe_writer(self) -> None: """Test writers.""" bufferings: list[int | None] = [None, -1, 0, 1024 * 1024] batch_sizes: list[int | None] = [None, 1, 2, 3] for buffering, batch_size in itt.product(bufferings, batch_sizes): with ( self.subTest(batch_size=batch_size, buffering=buffering), tempfile.TemporaryDirectory() as directory, ): path = Path(directory) / "test.tsv" with safe_open_writer(path, buffering=buffering, batch_size=batch_size) as writer: writer.writerow(("c1", "c2")) writer.writerow(("v1", "v2")) writer.writerow(("v3", "v4")) self.assertEqual( "c1\tc2\nv1\tv2\nv3\tv4\n", path.read_text(), msg=f"incorrectly written with {buffering=} and {batch_size=}", ) df = pd.read_csv(path, sep="\t") self.assertEqual(["c1", "c2"], list(df.columns)) self.assertEqual(2, len(df.index)) with safe_open_reader(path) as reader: self.assertEqual([["c1", "c2"], ["v1", "v2"], ["v3", "v4"]], list(reader)) with safe_open_dict_reader(path) as reader2: self.assertEqual( [{"c1": "v1", "c2": "v2"}, {"c1": "v3", "c2": "v4"}], list(reader2) ) def test_zip_writer(self) -> None: """Test ZIP writers.""" with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "test.zip" inner = "test_inner.tsv" with open_zip_writer(path, inner) as writer: writer.writerow(("c1", "c2")) writer.writerow(("v1", "v2")) writer.writerow(("v3", "v4")) df = read_zipfile_csv(path, inner) self.assertEqual(["c1", "c2"], list(df.columns)) self.assertEqual(2, len(df.index)) with open_zip_reader(path, inner) as reader: self.assertEqual([["c1", "c2"], ["v1", "v2"], ["v3", "v4"]], list(reader)) def test_zip_writer_exc(self) -> None: """Test throwing an exception.""" with self.assertRaises(ValueError), tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.zip") with open_zipfile(path, "test.tsv", operation="write", representation="lolno"): # type:ignore pass def test_tar_ope(self) -> None: """Test writing and reading a tar file.""" inner = "test_inner.tsv" with self.subTest(representation="binary"), tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.tar.gz") # failure on invalid operation with self.assertRaises(KeyError), open_tarfile(path, inner, operation="nope"): # type:ignore[call-overload] pass with open_tarfile(path, inner, operation="write", representation="binary") as file: file.write(b"c1\tc2\nv1\tv2") with open_tarfile(path, inner, operation="read", representation="binary") as file: self.assertEqual(b"c1\tc2\n", next(file)) self.assertEqual(b"v1\tv2", next(file)) with self.subTest(representation="text"), tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.tar.gz") with open_tarfile(path, inner, operation="write", representation="text") as file: file.write("c1\tc2\nv1\tv2") with open_tarfile(path, inner, operation="read", representation="text") as file: self.assertEqual("c1\tc2\n", next(file)) self.assertEqual("v1\tv2", next(file)) def test_open_inner_zip(self) -> None: """Test opening a file within a zip archive.""" inner_path = "test-1.csv" data = "c1,c2\nv1,v2" with tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.zip") with zipfile.ZipFile(path, "w") as zip_file: zip_file.writestr(inner_path, data) # default with utils.open_inner_zipfile(zip_file, inner_path=inner_path) as file: self.assertEqual(data, file.read()) # explicit text mode with utils.open_inner_zipfile( zip_file, inner_path=inner_path, representation="text" ) as file: self.assertEqual("TextIOWrapper", file.__class__.__name__) self.assertEqual(data, file.read()) with utils.open_inner_zipfile( zip_file, inner_path=inner_path, representation="binary" ) as file: self.assertEqual("ZipExtFile", file.__class__.__name__) self.assertEqual(data, file.read().decode("utf-8")) def test_zip_csvs(self) -> None: """Test reading many CSVs from inside a zip file.""" with tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.zip") with zipfile.ZipFile(path, "w") as zip_file: zip_file.writestr("test-1.csv", "c1,c2\nv1,v2") zip_file.writestr("test-2.csv", "c1,c2\nv3,v4") # this is a decoy file zip_file.writestr("test-3.xxx", "c1,c2\nv5,v6") rows = list(iter_zipped_csvs(path, progress=False)) self.assertEqual([["v1", "v2"], ["v3", "v4"]], rows) records = list(iter_zipped_csvs(path, progress=False, return_type="record")) self.assertEqual([{"c1": "v1", "c2": "v2"}, {"c1": "v3", "c2": "v4"}], records) def test_zip_csvs_header_mismatch(self) -> None: """Test raising an exception when the headers don't match.""" with tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.zip") with zipfile.ZipFile(path, "w") as zip_file: zip_file.writestr("test-1.csv", "c1,c2\nv1,v2") zip_file.writestr("test-2.csv", "c3,c4\nv3,v4") with self.assertRaises(ValueError): list(iter_zipped_csvs(path, progress=False)) def test_tarred_csvs(self) -> None: """Test reading many CSVs from inside a tar file.""" with tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.tar.gz") with tarfile.open(path, "w") as tar_file: tarfile_writestr(tar_file, "test-1.csv", "c1,c2\nv1,v2") tarfile_writestr(tar_file, "test-2.csv", "c1,c2\nv3,v4") # decoy, with wrong extension tarfile_writestr(tar_file, "test-3.xxx", "c1,c2\nv5,v6") rows = list(iter_tarred_csvs(path, progress=False)) self.assertEqual([["v1", "v2"], ["v3", "v4"]], rows) records = list(iter_tarred_csvs(path, progress=False, return_type="record")) self.assertEqual([{"c1": "v1", "c2": "v2"}, {"c1": "v3", "c2": "v4"}], records) def test_tarred_csvs_mismatch(self) -> None: """Test reading many CSVs from inside a tar file.""" with tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.tar.gz") with tarfile.open(path, "w") as tar_file: tarfile_writestr(tar_file, "test-1.csv", "c1,c2\nv1,v2") tarfile_writestr(tar_file, "test-2.csv", "c3,c4\nv3,v4") with self.assertRaises(ValueError): list(iter_tarred_csvs(path, progress=False)) def test_lzma(self) -> None: """Test LZMA.""" with tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.lzma") df = pd.DataFrame([("v1", "v2"), ("v3", "v4")], columns=["c1", "c2"]) write_lzma_csv(df, path) new_df = read_lzma_csv(path) self.assertEqual(new_df.values.tolist(), df.values.tolist()) def test_safe_open_exceptions(self) -> None: """Test exceptions raised by :func:`safe_open`.""" with self.assertRaises(ValueError), safe_open(TEST_TXT, representation="nope") as _file: # type:ignore pass with self.assertRaises(ValueError), safe_open(TEST_TXT, operation="nope") as _file: # type:ignore pass with self.assertRaises(TypeError), safe_open(5) as _file: # type:ignore pass with ( self.assertRaises(ValueError), safe_open(TEST_TXT, representation="binary", encoding="utf-8") as _file, ): pass with ( self.assertRaises(ValueError), safe_open(TEST_TXT, representation="binary", newline="") as _file_binary, ): pass url = "https://zenodo.org/records/15504009/files/startup.sh" with self.assertRaises(ValueError), safe_open(url, operation="write") as _file: pass def test_safe_open_binary(self) -> None: """Test safe open in binary mode.""" for path in TEST_PATHS: with self.subTest(path=path): with safe_open(path, representation="binary") as file: self.assertEqual( TEST_TXT_CONTENT, file.read().decode("utf-8"), msg=f"failed to read bytes from {path}", ) with ( safe_open(path, representation="binary") as passthrough, safe_open(passthrough, representation="binary") as file, ): self.assertEqual( TEST_TXT_CONTENT, file.read().decode("utf-8"), msg=f"failed to read bytes from {path} in a passthrough scenario", ) with ( safe_open(path, representation="binary") as passthrough, safe_open(passthrough, representation="text") as file, ): self.assertEqual( TEST_TXT_CONTENT, file.read(), msg=f"failed to wrap bytes from {path} in a passthrough scenario", ) def test_safe_open_url_binary(self) -> None: """Test safe open in URL mode.""" with safe_open( "https://zenodo.org/records/15504009/files/startup.sh", representation="binary" ) as file: self.assertIn("sleep 5", file.read().decode("utf-8")) def test_safe_open_text(self) -> None: """Test safe open in text mode.""" for path, encoding, newline in itt.product(TEST_PATHS, [None, "utf-8"], [None, "\n"]): with self.subTest(path=path, encoding=encoding, newline=newline): with safe_open( path, encoding=encoding, representation="text", newline=newline ) as file: self.assertEqual( TEST_TXT_CONTENT, file.read(), msg=f"failed to read text from {path}" ) with ( safe_open( path, encoding=encoding, representation="text", newline=newline ) as passthrough, safe_open(passthrough) as file, ): self.assertEqual( TEST_TXT_CONTENT, file.read(), msg=f"failed to read text from {path} in a passthrough scenario", ) with ( safe_open( path, encoding=encoding, representation="text", newline=newline ) as passthrough, safe_open(passthrough, representation="binary") as file, ): self.assertEqual( TEST_TXT_CONTENT, file.read().decode("utf-8"), msg=f"failed to buffer text from {path} in a passthrough scenario", ) def test_safe_open_url_text(self) -> None: """Test safe open in URL mode.""" with safe_open( "https://zenodo.org/records/15504009/files/startup.sh", representation="text" ) as file: self.assertIn("sleep 5", file.read()) def test_encodings(self) -> None: """Test I/O in different encodings.""" encodings = [None, "ascii", "utf-16-be", "CP1252"] compressions = [None, *COMPRESSION_EXTENSIONS] bufferings = [None, 1024 * 1024] for encoding, compression, buffering in itt.product(encodings, compressions, bufferings): with ( self.subTest(encoding=encoding, compression=compression, buffering=buffering), tempfile.TemporaryDirectory() as directory_name, ): directory = Path(directory_name) if compression is None: path = directory.joinpath("test.txt") else: path = directory.joinpath(f"test.txt.{compression}") with safe_open( path, encoding=encoding, operation="write", buffering=buffering ) as file: file.write(TEST_TXT_CONTENT) with safe_open(path, encoding=encoding, operation="read") as file: self.assertEqual(TEST_TXT_CONTENT, file.read()) def test_safe_write(self) -> None: """Test safe write.""" with tempfile.TemporaryDirectory() as directory: path = Path(directory).joinpath("test.txt") safe_write_text(TEST_TXT_CONTENT, path) self.assertEqual(TEST_TXT_CONTENT, safe_read_text(path)) with path.open() as file: self.assertEqual(TEST_TXT_CONTENT, file.read()) def test_gzip(self) -> None: """Test gzipping.""" with tempfile.TemporaryDirectory() as tmpdir: directory = Path(tmpdir) path = directory.joinpath("test.txt") path_gz = directory.joinpath("test.txt.gz") path.write_text(TEST_TXT_CONTENT) self.assertTrue(path.is_file()) self.assertFalse(path_gz.is_file()) gzip_compress(path, cleanup=False) self.assertTrue(path.is_file()) self.assertTrue(path_gz.is_file()) path_gz.unlink() gzip_compress(path, cleanup=True) self.assertFalse(path.is_file()) self.assertTrue(path_gz.is_file()) path_new = directory.joinpath("test-new.txt") gunzip(path_gz, path_new, cleanup=False) self.assertTrue(path_gz.is_file()) self.assertTrue(path_new.is_file()) self.assertEqual(TEST_TXT_CONTENT, path_new.read_text(encoding="utf-8")) path_new.unlink() gunzip(path_gz, path_new, cleanup=True) self.assertFalse(path_gz.is_file()) self.assertTrue(path_new.is_file()) self.assertEqual(TEST_TXT_CONTENT, path_new.read_text(encoding="utf-8")) def test_is_url(self) -> None: """Test checking URL.""" self.assertTrue(is_url("https://zenodo.org/records/15504009/files/startup.sh?download=1")) self.assertTrue(is_url("https://zenodo.org/records/15504009/files/startup.sh")) self.assertTrue(is_url("http://zenodo.org/records/15504009/files/startup.sh")) self.assertFalse(is_url("ftp://zenodo.org/records/15504009/files/startup.sh")) self.assertFalse(is_url("nope")) self.assertFalse(is_url(Path().cwd())) def test_open_url(self) -> None: """Test opening a URL.""" with open_url( "https://zenodo.org/records/15504009/files/startup.sh", representation="text", timeout=5 ) as file: self.assertIn("sleep 5", file.read()) with open_url( "https://zenodo.org/records/15504009/files/startup.sh", representation="binary", timeout=5, ) as file: self.assertIn("sleep 5", file.read().decode("utf-8")) def test_get_json(self) -> None: """Test getting JSON.""" url = "https://zenodo.org/records/15826754/files/configuration.json?download=1" x = safe_open_json(url, timeout=5) self.assertIsInstance(x, dict) self.assertIn("name", x) def test_safe_open_yaml(self) -> None: """Test getting YAML.""" url = "https://github.com/biopragmatics/mapping-registry/raw/refs/heads/main/registry.yml" x = safe_open_yaml(url, timeout=5) self.assertIsInstance(x, dict) self.assertIn("mapping_registry_id", x) class TestDownload(unittest.TestCase): """Tests for downloading.""" def setUp(self) -> None: """Set up a test.""" self.directory_obj = tempfile.TemporaryDirectory() self.directory = Path(self.directory_obj.name) self.bad_url = "https://nope.nope/nope.tsv" self.path_for_bad_url = self.directory.joinpath("nope.tsv") def tearDown(self) -> None: """Tear down a test.""" self.directory_obj.cleanup() def test_bad_file_error(self) -> None: """Test that urllib errors are handled properly.""" with self.assertRaises(DownloadError): download( url=self.bad_url, path=self.path_for_bad_url, backend="urllib", ) self.assertFalse(self.path_for_bad_url.is_file()) def test_requests_error_stream(self) -> None: """Test that requests errors are handled properly.""" with self.assertRaises(DownloadError): download( url=self.bad_url, path=self.path_for_bad_url, backend="requests", stream=True, ) self.assertFalse(self.path_for_bad_url.is_file()) def test_requests_error_sync(self) -> None: """Test that requests errors are handled properly.""" with self.assertRaises(DownloadError): download( url=self.bad_url, path=self.path_for_bad_url, backend="requests", stream=False, ) self.assertFalse(self.path_for_bad_url.is_file()) class TestHashing(unittest.TestCase): """Tests for hexdigest checking.""" def setUp(self) -> None: """Set up a test.""" self.directory = tempfile.TemporaryDirectory() self.path = Path(self.directory.name).joinpath("test.tsv") md5 = hashlib.md5() # noqa: S324 with TEST_TXT.open("rb") as file: md5.update(file.read()) self.expected_md5 = md5.hexdigest() self.mismatching_md5_hexdigest = "yolo" self.assertNotEqual(self.mismatching_md5_hexdigest, self.expected_md5) def tearDown(self) -> None: """Tear down a test.""" self.directory.cleanup() def test_hash_success(self) -> None: """Test checking actually works.""" self.assertFalse(self.path.exists()) download( url=TEST_TXT.as_uri(), path=self.path, hexdigests={ "md5": self.expected_md5, }, ) def test_hash_remote_success(self) -> None: """Test checking actually works.""" self.assertFalse(self.path.exists()) download( url=TEST_TXT.as_uri(), path=self.path, hexdigests_remote={ "md5": TEST_TXT_MD5.as_uri(), }, hexdigests_strict=True, ) self.assertTrue(self.path.exists()) def test_hash_remote_verbose_success(self) -> None: """Test checking actually works.""" self.assertFalse(self.path.exists()) download( url=TEST_TXT.as_uri(), path=self.path, hexdigests_remote={ "md5": TEST_TXT_VERBOSE_MD5.as_uri(), }, hexdigests_strict=False, ) self.assertTrue(self.path.exists()) def test_hash_remote_verbose_failure(self) -> None: """Test checking actually works.""" self.assertFalse(self.path.exists()) with self.assertRaises(HexDigestError): download( url=TEST_TXT.as_uri(), path=self.path, hexdigests_remote={ "md5": TEST_TXT_VERBOSE_MD5.as_uri(), }, hexdigests_strict=True, ) def test_hash_error(self) -> None: """Test hash error on download.""" self.assertFalse(self.path.exists()) with self.assertRaises(HexDigestError): download( url=TEST_TXT.as_uri(), path=self.path, hexdigests={ "md5": self.mismatching_md5_hexdigest, }, ) def test_hash_remote_error(self) -> None: """Test hash error on download.""" self.assertFalse(self.path.exists()) with self.assertRaises(HexDigestError): download( url=TEST_TXT.as_uri(), path=self.path, hexdigests_remote={ "md5": TEST_TXT_WRONG_MD5.as_uri(), }, hexdigests_strict=True, ) def test_override_hash_error(self) -> None: """Test hash error on download.""" self.path.write_text("test file content") self.assertTrue(self.path.exists()) with self.assertRaises(HexDigestError): download( url=TEST_TXT.as_uri(), path=self.path, hexdigests={ "md5": self.expected_md5, }, force=False, ) def test_override_hash_remote_error(self) -> None: """Test hash error on download.""" self.path.write_text("test file content") self.assertTrue(self.path.exists()) with self.assertRaises(HexDigestError): download( url=TEST_TXT.as_uri(), path=self.path, hexdigests_remote={ "md5": TEST_TXT_MD5.as_uri(), }, hexdigests_strict=True, force=False, ) def test_force(self) -> None: """Test overwriting wrong file.""" # now if force=True it should not bother with the hash check self.path.write_text("test file content") self.assertTrue(self.path.exists()) download( url=TEST_TXT.as_uri(), path=self.path, hexdigests={ "md5": self.expected_md5, }, force=True, ) def test_remote_force(self) -> None: """Test overwriting wrong file.""" # now if force=True it should not bother with the hash check self.path.write_text("test file content") self.assertTrue(self.path.exists()) download( url=TEST_TXT.as_uri(), path=self.path, hexdigests_remote={ "md5": TEST_TXT_MD5.as_uri(), }, hexdigests_strict=True, force=True, ) def test_hexdigest_urls(self) -> None: """Test getting hex digests from URLs.""" for url, strict in [ (TEST_TXT_MD5, True), (TEST_TXT_MD5, False), (TEST_TXT_VERBOSE_MD5, False), ]: hexdigests = get_hexdigests_remote( {"md5": url.as_uri()}, hexdigests_strict=strict, ) self.assertEqual( "4221d002ceb5d3c9e9137e495ceaa647", hexdigests["md5"], ) hexdigests = get_hexdigests_remote( {"md5": TEST_TXT_VERBOSE_MD5.as_uri()}, hexdigests_strict=True ) self.assertNotEqual( "4221d002ceb5d3c9e9137e495ceaa647", hexdigests["md5"], ) @unittest.skip(reason="this test hits a live endpoint") def test_live(self) -> None: """Test live.""" hexdigests = get_hexdigests_remote( {"md5": "https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed22n0001.xml.gz.md5"}, hexdigests_strict=False, ) self.assertEqual( { "md5": "0f08d8f3947dde1f3bced5e1f450c0da", }, hexdigests, ) class TestGitHub(unittest.TestCase): """Tests for GitHub.""" def test_repository(self) -> None: """Test getting a repository.""" data = get_repository("cthoyt", "pystow").json() self.assertEqual(318194121, data["id"]) self.assertEqual("main", data["default_branch"]) def test_default_branch(self) -> None: """Test getting the default branch.""" self.assertEqual("main", get_default_branch("cthoyt", "pystow")) cthoyt-pystow-54a586e/tox.ini000066400000000000000000000211131525577235500162630ustar00rootroot00000000000000# Tox (http://tox.testrun.org/) is a tool for running tests # in multiple virtualenvs. This configuration file will run the # test suite on all supported python versions. To use it, "pip install tox" # and then run "tox" from this directory. [tox] # To use a PEP 517 build-backend you are required to configure tox to use an isolated_build: # https://tox.readthedocs.io/en/latest/example/package.html isolated_build = True # These environments are run in order if you just use `tox`: envlist = # always keep coverage-clean first coverage-clean # code formatters format # format-docs # Code quality assessment pyroma lint mypy # Documentation quality assurance doc8 docstr-coverage docs-test # the actual tests py doctests # always keep coverage-report last coverage-report [testenv] description = Run unit and integration tests. # Runs on the "tests" directory by default, or passes the positional # arguments from `tox -e py ... commands = coverage run -p -m pytest --durations=20 {posargs:tests} coverage combine coverage xml dependency_groups = tests extras = # See the [project.optional-dependencies] entry in pyproject.toml pandas rdf xml ratelimit bs4 yaml pydantic requests cli [testenv:coverage-clean] description = Remove testing coverage artifacts. deps = coverage skip_install = true commands = coverage erase [testenv:doctests] description = Test that documentation examples run properly. commands = # note that the package name is required for discovery xdoctest -m src/pystow deps = xdoctest pygments extras = pandas yaml pydantic [testenv:treon] description = Test that notebooks can run to completion commands = treon notebooks/ deps = treon [testenv:format] description = Format the code in a deterministic way using ruff. Note that ruff check should come before ruff format when using --fix (ref: https://github.com/astral-sh/ruff-pre-commit/blob/main/README.md) deps = ruff skip_install = true commands = ruff check --fix ruff format [testenv:format-docs] description = Run documentation linters. # note that this doesn't work with sphinx-click # or any other extension that adds extra directives deps = rstfmt extras = # See the [project.optional-dependencies] entry in pyproject.toml for "docs" docs skip_install = true commands = rstfmt docs/source/ [testenv:format-markdown] description = Run markdown formatter. skip_install = true allowlist_externals = npx commands = npx --yes prettier --write --prose-wrap always "**/*.md" [testenv:lint] description = Check code quality using ruff and other tools. skip_install = true deps = ruff commands = ruff check ruff format --check [testenv:pyroma] deps = uv-build pygments pyroma skip_install = true commands = pyroma --min=10 . description = Run the pyroma tool to check the package friendliness of the project. [testenv:mypy] description = Run the mypy tool to check static typing on the project. Installs the package to make sure all type stubs get recognized. dependency_groups = typing extras = # See the [project.optional-dependencies] entry in pyproject.toml pandas rdf xml ratelimit bs4 yaml pydantic requests cli commands = mypy --install-types --non-interactive --ignore-missing-imports --strict src/ tests/ [testenv:doc8] skip_install = true deps = doc8 extras = docs commands = doc8 docs/source/ description = Run the doc8 tool to check the style of the RST files in the project docs. [testenv:docstr-coverage] description = Run the docstr-coverage tool to check documentation coverage. skip_install = true deps = docstr-coverage commands = docstr-coverage src/ tests/ --skip-private --skip-magic [testenv:docs] description = Build the documentation locally, allowing warnings. dependency_groups = docs commands = python -m sphinx -b html -d docs/build/doctrees docs/source docs/build/html [testenv:docs-test] description = Test building the documentation in an isolated environment. Warnings are considered as errors via -W. changedir = docs dependency_groups = {[testenv:docs]dependency_groups} commands = mkdir -p {envtmpdir} cp -r source {envtmpdir}/source python -m sphinx -W -b html -d {envtmpdir}/build/doctrees {envtmpdir}/source {envtmpdir}/build/html # python -m sphinx -W -b coverage -d {envtmpdir}/build/doctrees {envtmpdir}/source {envtmpdir}/build/coverage # cat {envtmpdir}/build/coverage/c.txt # cat {envtmpdir}/build/coverage/python.txt allowlist_externals = cp cat mkdir [testenv:coverage-xml] deps = coverage[toml] skip_install = true commands = coverage xml [testenv:coverage-report] # TODO this is broken deps = coverage[toml] skip_install = true commands = coverage report #################### # Deployment tools # #################### [testenv:bumpversion] description = Bump the version number commands = bump-my-version bump {posargs} skip_install = true passenv = HOME deps = bump-my-version [testenv:bumpversion-release] description = Remove the -dev tag from the version commands = bump-my-version bump release --tag skip_install = true passenv = HOME deps = bump-my-version [testenv:build] skip_install = true deps = uv commands = uv build --sdist --wheel --no-build-isolation ############ # Releases # ############ # In order to make a release to PyPI, you'll need to take the following steps: # # 1. Navigate to https://pypi.org/account/register/ to register for Test PyPI # 2. Navigate to https://pypi.org/manage/account/ and request to re-send a verification email. # This is not sent by default, and is required to set up 2-Factor Authentication. # 3. Get account recovery codes # 4. Set up 2-Factor Authentication # 5. Get an API token from https://pypi.org/manage/account/token/ # 6. Install keyring with `uv tool install keyring` # 7. Add your token to keyring with `keyring set https://upload.pypi.org/legacy/ __token__` [testenv:release] description = Release the code to PyPI so users can pip install it, using credentials from keyring skip_install = true deps = {[testenv:build]deps} uv keyring commands = {[testenv:build]commands} uv publish --username __token__ --keyring-provider subprocess --publish-url https://upload.pypi.org/legacy/ [testenv:release-via-env] description = Release the code to PyPI so users can pip install it, using credentials from the environment. skip_install = true deps = {[testenv:build]deps} uv commands = {[testenv:build]commands} uv publish --publish-url https://upload.pypi.org/legacy/ passenv = UV_PUBLISH_USERNAME UV_PUBLISH_PASSWORD [testenv:finish] description = Run a workflow that removes -dev from the version, creates a tagged release on GitHub, creates a release on PyPI, and bumps the version again. skip_install = true passenv = HOME deps = {[testenv:release]deps} bump-my-version commands = {[testenv:bumpversion-release]commands} {[testenv:release]commands} git push --tags bump-my-version bump patch git push allowlist_externals = git ################# # Test Releases # ################# # In order to test making a release to Test PyPI, you'll need to take the following steps: # # 1. Navigate to https://test.pypi.org/account/register/ to register for Test PyPI # 2. Navigate to https://test.pypi.org/manage/account/ and request to re-send a verification email. # This is not sent by default, and is required to set up 2-Factor Authentication. # 3. Get account recovery codes # 4. Set up 2-Factor Authentication # 5. Get an API token from https://test.pypi.org/manage/account/token/ # 6. Install keyring with `uv tool install keyring` # 7. Add your token to keyring with `keyring set https://test.pypi.org/legacy/ __token__` [testenv:testrelease] description = Release the code to the test PyPI site skip_install = true deps = {[testenv:build]deps} uv keyring commands = {[testenv:build]commands} uv publish --username __token__ --keyring-provider subprocess --publish-url https://test.pypi.org/legacy/ [testenv:testfinish] description = Run a workflow that removes -dev from the version, creates a tagged release on GitHub, creates a release on Test PyPI, and bumps the version again. skip_install = true passenv = HOME deps = {[testenv:testrelease]deps} bump-my-version commands = {[testenv:bumpversion-release]commands} {[testenv:testrelease]commands} git push --tags bump-my-version bump patch git push allowlist_externals = git