├── compute_modules ├── py.typed ├── bin │ ├── __init__.py │ ├── ontology │ │ ├── __init__.py │ │ ├── _config_path.py │ │ ├── _types.py │ │ ├── metadata_client.py │ │ ├── generate_metadata_config.py │ │ └── metadata_loader.py │ └── static_inference │ │ ├── __init__.py │ │ └── infer.py ├── client │ ├── __init__.py │ └── encoder.py ├── function_registry │ ├── __init__.py │ ├── datetime_conversion_util.py │ ├── function.py │ ├── types.py │ ├── function_payload_converter.py │ └── function_registry.py ├── sources_v2 │ ├── __init__.py │ ├── _back_compat.py │ ├── _api.py │ └── _sources.py ├── _version.py ├── arguments │ ├── __init__.py │ └── arguments.py ├── resources │ ├── types.py │ ├── __init__.py │ └── pipeline_resources.py ├── context │ ├── __init__.py │ ├── context.py │ └── types.py ├── logging │ ├── __init__.py │ ├── public.py │ ├── internal.py │ └── common.py ├── auth │ ├── __init__.py │ ├── pipeline.py │ └── third_party.py ├── __init__.py ├── sources │ ├── __init__.py │ └── _sources.py ├── annotations.py └── startup.py ├── changelog ├── @unreleased │ ├── .gitkeep │ └── pr-40.v2.yml ├── 0.4.0 │ ├── pr-6.v2.yml │ ├── pr-7.v2.yml │ └── pr-5.v2.yml ├── 0.23.0 │ ├── pr-34.v2.yml │ └── pr-35.v2.yml ├── 0.8.0 │ └── pr-12.v2.yml ├── 0.17.0 │ ├── pr-14.v2.yml │ └── pr-24.v2.yml ├── 0.19.0 │ └── pr-29.v2.yml ├── 0.20.0 │ ├── pr-30.v2.yml │ └── pr-31.v2.yml ├── 0.5.0 │ ├── pr-8.v2.yml │ └── pr-9.v2.yml ├── 0.15.0 │ └── pr-23.v2.yml ├── 0.2.0 │ ├── pr-1.v2.yml │ └── pr-2.v2.yml ├── 0.25.0 │ └── pr-37.v2.yml ├── 0.6.0 │ └── pr-10.v2.yml ├── 0.7.0 │ └── pr-11.v2.yml ├── 0.13.0 │ └── pr-21.v2.yml ├── 0.16.0 │ └── pr-25.v2.yml ├── 0.3.0 │ └── pr-4.v2.yml ├── 0.9.0 │ └── pr-15.v2.yml ├── 0.10.0 │ └── pr-18.v2.yml ├── 0.18.0 │ └── pr-28.v2.yml ├── 0.12.0 │ └── pr-19.v2.yml ├── 0.21.0 │ ├── pr-24.v2.yml │ └── pr-26.v2.yml ├── 0.22.0 │ └── pr-32.v2.yml ├── 0.11.0 │ └── pr-17.v2.yml ├── 0.14.0 │ └── pr-20.v2.yml ├── 0.24.0 │ └── pr-38.v2.yml └── 0.26.0 │ └── pr-39.v2.yml ├── assets └── arguments_example.png ├── .changelog.yml ├── .github ├── ISSUE_TEMPLATE.md ├── PULL_REQUEST_TEMPLATE.md └── stale.yml ├── .bulldozer.yml ├── .excavator.yml ├── LICENSE.txt ├── scripts ├── __init__.py ├── ontology │ ├── __init__.py │ ├── _config_path.py │ ├── generate_metadata_config.py │ ├── metadata_client.py │ ├── main.py │ ├── _types.py │ └── metadata_loader.py ├── set_version.py └── checks.py ├── tests ├── __init__.py ├── infer │ ├── __init__.py │ ├── test_project │ │ ├── __init__.py │ │ ├── decorated │ │ │ ├── __init__.py │ │ │ ├── collection_types.py │ │ │ ├── alias_types.py │ │ │ └── primitive_types.py │ │ ├── manually_registered │ │ │ ├── __init__.py │ │ │ ├── unguarded_single_function.py │ │ │ ├── single_function.py │ │ │ ├── mixed_registration.py │ │ │ └── multiple_functions.py │ │ └── _types.py │ └── test_infer.py ├── logging │ ├── __init__.py │ ├── logging_test_utils.py │ ├── test_log_levels.py │ └── test_log_context.py ├── function_registry │ ├── __init__.py │ ├── dummy_app_with_issues.py │ ├── dummy_app.py │ ├── test_function_payload_converter.py │ └── test_function_schema_parser.py ├── conftest.py └── auth │ └── test_third_party.py ├── DEV_SETUP.md ├── pyproject.toml ├── .circleci └── config.yml └── .gitignore /compute_modules/py.typed: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /changelog/@unreleased/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /assets/arguments_example.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/palantir/python-compute-module/HEAD/assets/arguments_example.png -------------------------------------------------------------------------------- /changelog/0.4.0/pr-6.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: Fix ci checks 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/6 6 | -------------------------------------------------------------------------------- /changelog/0.23.0/pr-34.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: Bump 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/34 6 | -------------------------------------------------------------------------------- /changelog/0.8.0/pr-12.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: set_version error poetry cmd 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/12 6 | -------------------------------------------------------------------------------- /changelog/0.17.0/pr-14.v2.yml: -------------------------------------------------------------------------------- 1 | type: manualTask 2 | manualTask: 3 | description: Add poe task runner 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/14 6 | -------------------------------------------------------------------------------- /changelog/0.19.0/pr-29.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: Fix missing kwargs in query context 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/29 6 | -------------------------------------------------------------------------------- /changelog/0.20.0/pr-30.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: Add refreshing oauth token 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/30 6 | -------------------------------------------------------------------------------- /changelog/0.5.0/pr-8.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: adding @function to docs 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/8 6 | -------------------------------------------------------------------------------- /changelog/0.15.0/pr-23.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: ThreadPool concurrency support 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/23 6 | -------------------------------------------------------------------------------- /changelog/0.2.0/pr-1.v2.yml: -------------------------------------------------------------------------------- 1 | type: manualTask 2 | manualTask: 3 | description: Enforce license in all files 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/1 6 | -------------------------------------------------------------------------------- /changelog/0.25.0/pr-37.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: Fix telemetry SLS parsing 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/37 6 | -------------------------------------------------------------------------------- /changelog/0.4.0/pr-7.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: Raise error early on dict with no type params 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/7 6 | -------------------------------------------------------------------------------- /changelog/0.6.0/pr-10.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: force encoding to be utf-8 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/10 6 | -------------------------------------------------------------------------------- /changelog/0.7.0/pr-11.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: Remove executable scripts 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/11 6 | -------------------------------------------------------------------------------- /changelog/0.13.0/pr-21.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: Send OAuth scopes as space-separated string 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/21 6 | -------------------------------------------------------------------------------- /changelog/0.16.0/pr-25.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: Get optional uerId from Job 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/25 6 | -------------------------------------------------------------------------------- /changelog/0.3.0/pr-4.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: Add static typing for QueryContext param 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/4 6 | -------------------------------------------------------------------------------- /changelog/0.9.0/pr-15.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: detect stream result and post 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/15 6 | -------------------------------------------------------------------------------- /changelog/0.2.0/pr-2.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: Include job ID & process ID in logger context 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/2 6 | -------------------------------------------------------------------------------- /changelog/0.20.0/pr-31.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: fix silent error raised when executing unknown query 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/31 6 | -------------------------------------------------------------------------------- /changelog/0.4.0/pr-5.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: Handle Unknown output datatype with string as default 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/5 6 | -------------------------------------------------------------------------------- /changelog/0.10.0/pr-18.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: log error_body when post result fails 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/18 6 | -------------------------------------------------------------------------------- /changelog/0.18.0/pr-28.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: '[Connectivity] Support mounted source config' 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/28 6 | -------------------------------------------------------------------------------- /changelog/0.5.0/pr-9.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: Pypi package can be publish to conda-forge 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/9 6 | -------------------------------------------------------------------------------- /changelog/0.12.0/pr-19.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: Support default values for function inputs 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/19 6 | -------------------------------------------------------------------------------- /changelog/0.17.0/pr-24.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: add CLI for running schema inference at build time 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/24 6 | -------------------------------------------------------------------------------- /changelog/0.21.0/pr-24.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: add CLI for running schema inference at build time 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/24 6 | -------------------------------------------------------------------------------- /changelog/0.23.0/pr-35.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: Fix foundry-compute-modules[sources] package name in exception 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/35 6 | -------------------------------------------------------------------------------- /changelog/0.22.0/pr-32.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: '[Connectivity] Source interface in Compute Modules Lib' 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/32 6 | -------------------------------------------------------------------------------- /changelog/0.21.0/pr-26.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: Support for `ontologyProvenance` in static function schema generation 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/26 6 | -------------------------------------------------------------------------------- /changelog/0.11.0/pr-17.v2.yml: -------------------------------------------------------------------------------- 1 | type: improvement 2 | improvement: 3 | description: 'Performance improvements: process pooling; reusing connections' 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/17 6 | -------------------------------------------------------------------------------- /changelog/0.14.0/pr-20.v2.yml: -------------------------------------------------------------------------------- 1 | type: feature 2 | feature: 3 | description: Adding custom logging formatting support to native Compute Module Logging 4 | links: 5 | - https://github.com/palantir/python-compute-module/pull/20 6 | -------------------------------------------------------------------------------- /.changelog.yml: -------------------------------------------------------------------------------- 1 | # Excavator auto-updates this file. Please contribute improvements to the central template. 2 | 3 | # This file is intentionally empty. The file's existence enables changelog-app and is empty to use the default configuration. 4 | -------------------------------------------------------------------------------- /changelog/0.24.0/pr-38.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: Timestamps can now be in the format with fractional seconds through 4 | Python datetime as many timestamps are. 5 | links: 6 | - https://github.com/palantir/python-compute-module/pull/38 7 | -------------------------------------------------------------------------------- /changelog/0.26.0/pr-39.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: 'After this fix: If the job posting fails 5 times (e.g. result too 4 | large) then the client tries 5 more times to post just a simple error message.' 5 | links: 6 | - https://github.com/palantir/python-compute-module/pull/39 7 | -------------------------------------------------------------------------------- /changelog/@unreleased/pr-40.v2.yml: -------------------------------------------------------------------------------- 1 | type: fix 2 | fix: 3 | description: "Inform the 4 | forwarder whenever a node starts up so it will remove all existing jobs related 5 | to it that it thinks are still running." 6 | links: 7 | - https://github.com/palantir/python-compute-module/pull/40 8 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | ## What happened? 2 | 3 | 7 | 8 | ## What did you want to happen? 9 | 10 | 13 | -------------------------------------------------------------------------------- /.github/PULL_REQUEST_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | ## Before this PR 2 | 3 | 4 | ## After this PR 5 | 6 | ==COMMIT_MSG== 7 | ==COMMIT_MSG== 8 | 9 | ## Possible downsides? 10 | 11 | 12 | -------------------------------------------------------------------------------- /.github/stale.yml: -------------------------------------------------------------------------------- 1 | 2 | only: pulls 3 | staleLabel: stale 4 | pulls: 5 | daysUntilStale: 14 6 | daysUntilClose: 7 7 | exemptLabels: [ long-lived ] 8 | markComment: > 9 | This PR has been automatically marked as stale because it has not been touched in the last 14 days. 10 | If you'd like to keep it open, please leave a comment or add the 'long-lived' label, otherwise it'll be closed in 7 days. 11 | closeComment: false 12 | -------------------------------------------------------------------------------- /.bulldozer.yml: -------------------------------------------------------------------------------- 1 | # Excavator auto-updates this file. Please contribute improvements to the central template. 2 | 3 | version: 1 4 | merge: 5 | trigger: 6 | labels: ["merge when ready"] 7 | ignore: 8 | labels: ["do not merge"] 9 | method: squash 10 | options: 11 | squash: 12 | body: pull_request_body 13 | message_delimiter: ==COMMIT_MSG== 14 | delete_after_merge: true 15 | update: 16 | trigger: 17 | labels: ["update me"] 18 | -------------------------------------------------------------------------------- /.excavator.yml: -------------------------------------------------------------------------------- 1 | # Excavator auto-updates this file. Please contribute improvements to the central template. 2 | 3 | auto-label: 4 | names: 5 | versions-props/upgrade-all: [ "merge when ready" ] 6 | circleci/manage-circleci: [ "merge when ready" ] 7 | tags: 8 | donotmerge: [ "do not merge" ] 9 | roomba: [ "merge when ready", "🤖 fix nits" ] 10 | automerge: [ "merge when ready", "🤖 fix nits" ] 11 | standards: [ "merge when ready", "🤖 fix nits" ] 12 | autorelease: [ "autorelease" ] 13 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Copyright {YEAR} Palantir Technologies, Inc. 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | http://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. -------------------------------------------------------------------------------- /scripts/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/infer/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/logging/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /compute_modules/bin/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /scripts/ontology/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /compute_modules/client/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/function_registry/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/infer/test_project/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /compute_modules/bin/ontology/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /compute_modules/function_registry/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/infer/test_project/decorated/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /compute_modules/bin/static_inference/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/infer/test_project/manually_registered/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /compute_modules/sources_v2/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from ._sources import get_source 16 | 17 | __all__ = [ 18 | "get_source", 19 | ] 20 | -------------------------------------------------------------------------------- /compute_modules/_version.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | # The version is set during the publishing step (since we can't know the version in advance) 17 | # using the autorelease bot 18 | 19 | __version__ = "0.0.0" 20 | -------------------------------------------------------------------------------- /compute_modules/arguments/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from .arguments import get_parsed_arguments, get_raw_arguments 17 | 18 | __all__ = [ 19 | "get_parsed_arguments", 20 | "get_raw_arguments", 21 | ] 22 | -------------------------------------------------------------------------------- /compute_modules/resources/types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import typing 16 | from dataclasses import dataclass 17 | 18 | 19 | @dataclass 20 | class PipelineResource: 21 | rid: str 22 | branch: typing.Optional[str] = None 23 | -------------------------------------------------------------------------------- /compute_modules/context/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from .context import get_extra_context_parameters 17 | from .types import QueryContext 18 | 19 | __all__ = [ 20 | "get_extra_context_parameters", 21 | "QueryContext", 22 | ] 23 | -------------------------------------------------------------------------------- /compute_modules/resources/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from .pipeline_resources import get_pipeline_resources 16 | from .types import PipelineResource 17 | 18 | __all__ = [ 19 | "get_pipeline_resources", 20 | "PipelineResource", 21 | ] 22 | -------------------------------------------------------------------------------- /compute_modules/logging/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from .internal import set_internal_log_level 17 | from .public import get_logger, setup_logger_formatter 18 | 19 | __all__ = ["get_logger", "set_internal_log_level", "setup_logger_formatter"] 20 | -------------------------------------------------------------------------------- /compute_modules/auth/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from .pipeline import retrieve_pipeline_token 17 | from .third_party import RefreshingOauthToken, oauth, retrieve_third_party_id_and_creds 18 | 19 | __all__ = [ 20 | "oauth", 21 | "RefreshingOauthToken", 22 | "retrieve_pipeline_token", 23 | "retrieve_third_party_id_and_creds", 24 | ] 25 | -------------------------------------------------------------------------------- /compute_modules/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from ._version import __version__ as __version__ 17 | from .function_registry.function_registry import add_function, add_functions 18 | from .startup import ConcurrencyType, start_compute_module 19 | 20 | __all__ = [ 21 | "add_function", 22 | "add_functions", 23 | "ConcurrencyType", 24 | "start_compute_module", 25 | ] 26 | -------------------------------------------------------------------------------- /scripts/ontology/_config_path.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import os 17 | from typing import Optional 18 | 19 | DEFAULT_ONTOLOGY_METADATA_CONFIG_FILENAME = "ontology_metadata_config.json" 20 | 21 | 22 | def get_ontology_config_file(ontology_metadata_config_file: Optional[str]) -> str: 23 | if not ontology_metadata_config_file: 24 | ontology_metadata_config_file = os.path.join(os.getcwd(), DEFAULT_ONTOLOGY_METADATA_CONFIG_FILENAME) 25 | return ontology_metadata_config_file 26 | -------------------------------------------------------------------------------- /compute_modules/bin/ontology/_config_path.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import os 17 | from typing import Optional 18 | 19 | DEFAULT_ONTOLOGY_METADATA_CONFIG_FILENAME = "ontology_metadata_config.json" 20 | 21 | 22 | def get_ontology_config_file(ontology_metadata_config_file: Optional[str]) -> str: 23 | if not ontology_metadata_config_file: 24 | ontology_metadata_config_file = os.path.join(os.getcwd(), DEFAULT_ONTOLOGY_METADATA_CONFIG_FILENAME) 25 | return ontology_metadata_config_file 26 | -------------------------------------------------------------------------------- /compute_modules/auth/pipeline.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import os 17 | 18 | BUILD2_TOKEN = "BUILD2_TOKEN" 19 | 20 | 21 | def retrieve_pipeline_token() -> str: 22 | """Produces a bearer token that can be used to make calls to access pipeline resources. 23 | This is only available in pipeline mode. 24 | """ 25 | if BUILD2_TOKEN not in os.environ: 26 | raise RuntimeError("Pipeline token not available. Please make sure you are running in Pipeline mode.") 27 | with open(os.environ["BUILD2_TOKEN"], encoding="utf-8") as f: 28 | bearer_token = f.read() 29 | return bearer_token 30 | -------------------------------------------------------------------------------- /compute_modules/logging/public.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from logging import Formatter 16 | 17 | from .common import COMPUTE_MODULES_ADAPTER_MANAGER, ComputeModulesLoggerAdapter, _setup_logger_formatter 18 | 19 | 20 | def setup_logger_formatter(formatter: Formatter) -> None: 21 | _setup_logger_formatter(formatter) 22 | 23 | 24 | def get_logger(name: str) -> ComputeModulesLoggerAdapter: 25 | """Creates a logger instance for use within a compute module""" 26 | return COMPUTE_MODULES_ADAPTER_MANAGER.get_logger(name) 27 | 28 | 29 | __all__ = [ 30 | "get_logger", 31 | "setup_logger_formatter", 32 | ] 33 | -------------------------------------------------------------------------------- /tests/infer/test_project/manually_registered/unguarded_single_function.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from dataclasses import dataclass 16 | from typing import Dict 17 | 18 | from compute_modules.context import QueryContext 19 | from compute_modules.function_registry.function_registry import add_function 20 | from compute_modules.startup import start_compute_module 21 | 22 | 23 | @dataclass 24 | class DictWrapper: 25 | value: Dict[str, str] 26 | 27 | 28 | def return_dict_length_in_main(context: QueryContext, wrapper: DictWrapper) -> int: 29 | return len(wrapper.value) 30 | 31 | 32 | add_function(return_dict_length_in_main) 33 | start_compute_module() 34 | -------------------------------------------------------------------------------- /tests/infer/test_project/manually_registered/single_function.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from dataclasses import dataclass 16 | 17 | from compute_modules.context import QueryContext 18 | from compute_modules.function_registry.function_registry import add_function 19 | from compute_modules.startup import start_compute_module 20 | 21 | 22 | @dataclass 23 | class UndecoratedIntegerWrapper: 24 | value: int 25 | 26 | 27 | def return_integer_in_main(context: QueryContext, wrapper: UndecoratedIntegerWrapper) -> int: 28 | return wrapper.value 29 | 30 | 31 | if __name__ == "__main__": 32 | add_function(return_integer_in_main) 33 | start_compute_module() 34 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import json 16 | import logging 17 | from typing import Any 18 | 19 | import pytest 20 | 21 | 22 | class JsonFormatter(logging.Formatter): 23 | def format(self, record: Any) -> str: 24 | log_record = { 25 | "level": record.levelname, 26 | "process_id": record.process_id, 27 | "job_id": record.job_id, 28 | "location": f"{record.filename}:{record.lineno}", 29 | "message": record.getMessage(), 30 | "custom_text": "custom-message", 31 | } 32 | return json.dumps(log_record) 33 | 34 | 35 | @pytest.fixture 36 | def custom_formatter() -> JsonFormatter: 37 | return JsonFormatter() 38 | -------------------------------------------------------------------------------- /compute_modules/sources/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from ._sources import ( 16 | SOURCE_CONFIGURATIONS_PATH, 17 | SOURCE_CREDENTIALS_PATH, 18 | MountedHttpConnectionConfig, 19 | MountedSourceConfig, 20 | _source_configurations, 21 | _source_credentials, 22 | get_source_config, 23 | get_source_configurations, 24 | get_source_secret, 25 | get_sources, 26 | ) 27 | 28 | __all__ = [ 29 | "SOURCE_CONFIGURATIONS_PATH", 30 | "SOURCE_CREDENTIALS_PATH", 31 | "MountedHttpConnectionConfig", 32 | "MountedSourceConfig", 33 | "_source_configurations", 34 | "_source_credentials", 35 | "get_source_config", 36 | "get_source_configurations", 37 | "get_source_secret", 38 | "get_sources", 39 | ] 40 | -------------------------------------------------------------------------------- /scripts/set_version.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import subprocess 17 | 18 | VERSION_FILE_PATH = "compute_modules/_version.py" 19 | 20 | 21 | def _get_current_tag() -> str: 22 | return subprocess.check_output("git describe --tags --abbrev=0".split()).decode().strip() 23 | 24 | 25 | def main() -> None: 26 | gitversion = _get_current_tag() 27 | print(f"Setting {VERSION_FILE_PATH} to {gitversion}...") 28 | 29 | with open(VERSION_FILE_PATH, "r", encoding="utf-8") as f: 30 | content = f.read() 31 | 32 | content = content.replace('__version__ = "0.0.0"', f'__version__ = "{gitversion}"') 33 | 34 | with open(VERSION_FILE_PATH, "w", encoding="utf-8") as f: 35 | f.write(content) 36 | 37 | subprocess.run(["poetry", "version", gitversion]) 38 | print("Done!") 39 | -------------------------------------------------------------------------------- /compute_modules/context/context.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from typing import Any, Dict 17 | 18 | from ..auth import retrieve_third_party_id_and_creds 19 | from ..sources_v2._back_compat import get_mounted_source_secrets, get_mounted_sources 20 | 21 | 22 | def get_extra_context_parameters() -> Dict[str, Any]: 23 | formatted_source_configurations = {key: value.source_configuration for key, value in get_mounted_sources().items()} 24 | 25 | context_parameters: Dict[str, Any] = { 26 | "sources": get_mounted_source_secrets(), 27 | "source_configs": formatted_source_configurations, 28 | } 29 | 30 | CLIENT_ID, CLIENT_SECRET = retrieve_third_party_id_and_creds() 31 | 32 | if CLIENT_ID and CLIENT_SECRET: 33 | context_parameters.update({"CLIENT_ID": CLIENT_ID, "CLIENT_SECRET": CLIENT_SECRET}) 34 | 35 | return context_parameters 36 | -------------------------------------------------------------------------------- /tests/logging/logging_test_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from compute_modules.logging import internal 17 | 18 | DEBUG_STR = "I'm a little teapot" 19 | INFO_STR = "Short and stout" 20 | WARNING_STR = "Here is my handle, here is my spout" 21 | ERROR_STR = "When I get all steamed up hear me shout:" 22 | CRITICAL_STR = "Tip me over and pour me out!" 23 | 24 | CLIENT_DEBUG_STR = "twinkle twinkle little star" 25 | CLIENT_INFO_STR = "how I wonder what you are" 26 | CLIENT_WARNING_STR = "up above the world so high" 27 | CLIENT_ERROR_STR = "like a diamond in the sky" 28 | CLIENT_CRITICAL_STR = "oh nvm that's a planet" 29 | 30 | 31 | def emit_internal_logs() -> None: 32 | logger = internal.get_internal_logger() 33 | logger.debug(DEBUG_STR) 34 | logger.info(INFO_STR) 35 | logger.warning(WARNING_STR) 36 | logger.error(ERROR_STR) 37 | logger.critical(CRITICAL_STR) 38 | -------------------------------------------------------------------------------- /tests/infer/test_project/decorated/collection_types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from typing import Dict, List, Set 16 | 17 | from compute_modules.annotations import function 18 | from compute_modules.context import QueryContext 19 | from tests.infer.test_project._types import DummyOntologyType, OntologyEdit 20 | 21 | 22 | @function 23 | def return_list(context: QueryContext, event) -> List[str]: # type: ignore[no-untyped-def] 24 | return ["Hello", "World"] 25 | 26 | 27 | @function 28 | def return_set(context: QueryContext, event) -> Set[str]: # type: ignore[no-untyped-def] 29 | return {"Hello", "World"} 30 | 31 | 32 | @function 33 | def return_dict(context: QueryContext, event) -> Dict[str, str]: # type: ignore[no-untyped-def] 34 | return {"Hello": "World"} 35 | 36 | 37 | @function(edits=[DummyOntologyType]) 38 | def ontology_add_function(context: QueryContext, event) -> list[OntologyEdit]: # type: ignore[no-untyped-def] 39 | return [] 40 | -------------------------------------------------------------------------------- /compute_modules/annotations.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import atexit 16 | from typing import Any, Callable, List, Optional 17 | 18 | from .function_registry.function import Function 19 | from .function_registry.function_registry import add_function 20 | from .startup import start_compute_module 21 | 22 | 23 | def function( 24 | maybe_func: Any = None, 25 | *, 26 | streaming: bool = False, 27 | edits: Optional[List[Any]] = None, 28 | ) -> Callable[..., Any]: 29 | def function_wrapper(func: Callable[..., Any]) -> Callable[..., Any]: 30 | add_function(func, streaming=streaming, edits=edits) 31 | return Function(func, [] if edits is None else edits) 32 | 33 | if callable(maybe_func): 34 | return function_wrapper(maybe_func) 35 | return function_wrapper 36 | 37 | 38 | # Register the on_exit function to be called when the interpreter exits 39 | atexit.register(start_compute_module) 40 | 41 | __all__ = [ 42 | "function", 43 | ] 44 | -------------------------------------------------------------------------------- /compute_modules/logging/internal.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import logging 17 | from typing import Union 18 | 19 | from .common import COMPUTE_MODULES_ADAPTER_MANAGER, ComputeModulesLoggerAdapter 20 | 21 | INTERNAL_LOGGER_ADAPTER = None 22 | 23 | 24 | def set_internal_log_level(level: Union[str, int]) -> None: 25 | """Set the log level of the compute_modules_internal logger""" 26 | get_internal_logger().setLevel(level=level) 27 | 28 | 29 | def get_internal_logger() -> ComputeModulesLoggerAdapter: 30 | """Provides the internal ComputeModulesLoggerAdapter singleton""" 31 | global INTERNAL_LOGGER_ADAPTER 32 | if not INTERNAL_LOGGER_ADAPTER: 33 | INTERNAL_LOGGER_ADAPTER = COMPUTE_MODULES_ADAPTER_MANAGER.get_logger( 34 | "compute_modules_internal", 35 | default_level=logging.ERROR, 36 | ) 37 | return INTERNAL_LOGGER_ADAPTER 38 | 39 | 40 | __all__ = [ 41 | "get_internal_logger", 42 | "set_internal_log_level", 43 | ] 44 | -------------------------------------------------------------------------------- /tests/infer/test_project/_types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from typing import Any, NamedTuple, Optional, Union 17 | 18 | 19 | class DummyOntologyType: 20 | @staticmethod 21 | def api_name() -> str: 22 | return "DummyOntologyType" 23 | 24 | 25 | class ObjectLocator(NamedTuple): 26 | object_api_name: str 27 | primary_key: Any 28 | 29 | 30 | class AddObject(NamedTuple): 31 | locator: ObjectLocator 32 | fields: dict[str, Any] 33 | 34 | 35 | class ModifyObject(NamedTuple): 36 | locator: ObjectLocator 37 | fields: dict[str, Optional[Any]] 38 | 39 | 40 | class DeleteObject(NamedTuple): 41 | locator: ObjectLocator 42 | 43 | 44 | class AddLink(NamedTuple): 45 | relationship_api_name: str 46 | source: ObjectLocator 47 | target: ObjectLocator 48 | 49 | 50 | class RemoveLink(NamedTuple): 51 | relationship_api_name: str 52 | source: ObjectLocator 53 | target: ObjectLocator 54 | 55 | 56 | OntologyEdit = Union[AddObject, ModifyObject, DeleteObject, AddLink, RemoveLink] 57 | -------------------------------------------------------------------------------- /DEV_SETUP.md: -------------------------------------------------------------------------------- 1 | # `foundry-compute-modules` dev README 2 | 3 | [![Autorelease](https://img.shields.io/badge/Perform%20an-Autorelease-success.svg)](https://autorelease.general.dmz.palantir.tech/palantir/python-compute-module) 4 | 5 | ## Requirements 6 | 7 | * python >= 3.9 8 | * [poetry](https://python-poetry.org/docs/) 9 | 10 | ## Commands 11 | 12 | ### Install packages + scripts 13 | 14 | This command will install all deps specified in `pyproject.toml` and make the scripts specified in `pyproject.toml` available in your python environment. 15 | 16 | ```sh 17 | poetry install --extras sources 18 | ``` 19 | 20 | ### Run tests 21 | ```sh 22 | poe test 23 | ``` 24 | 25 | ### Run `mypy` checks 26 | ```sh 27 | poe check_mypy 28 | ``` 29 | 30 | ### Run linter checks (black, ruff, isort) 31 | This will run the same checks as those that are run during CI checks - meaning it will raise any issues found, but not fix them. 32 | 33 | ```sh 34 | poe check_format 35 | ``` 36 | 37 | ### Run formatter (black, ruff, isort) 38 | This will actually modify source files to fix any issues identified 39 | ```sh 40 | poe format 41 | ``` 42 | 43 | ### Build the library locally 44 | This will produce a `tar.gz` and a `whl` file in the `./dist` directory. 45 | 46 | ```sh 47 | poetry build 48 | ``` 49 | 50 | Either of these files can be installed locally for testing. For example: 51 | 52 | ```sh 53 | % poetry build 54 | Building foundry-compute-modules (0.0.0) 55 | - Building sdist 56 | - Built foundry_compute_modules-0.0.0.tar.gz 57 | - Building wheel 58 | - Built foundry_compute_modules-0.0.0-py3-none-any.whl 59 | 60 | % pip install ./dist/foundry_compute_modules-0.0.0.tar.gz 61 | ``` -------------------------------------------------------------------------------- /compute_modules/arguments/arguments.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from argparse import ArgumentParser, Namespace 17 | from sys import argv 18 | from typing import List 19 | 20 | 21 | def get_raw_arguments() -> List[str]: 22 | """Get command line arguments passed into the compute module as-is""" 23 | return argv[1:] 24 | 25 | 26 | def get_parsed_arguments() -> Namespace: 27 | """Get command line arguments passed into the compute module, parsed into an `argparse.Namespace` object. 28 | 29 | Note: for this to work properly you must pass args in a standard format, e.g. `--some-flag=test` or `-flag hello`. 30 | If you are using a non-standard format for args then use `get_raw_arguments` instead. 31 | """ 32 | # Source: https://stackoverflow.com/a/37367814 33 | parser = ArgumentParser() 34 | _, unknown = parser.parse_known_args() 35 | for arg in unknown: 36 | if arg.startswith(("-", "--")): 37 | parser.add_argument(arg.split("=")[0]) 38 | args = parser.parse_args() 39 | return args 40 | -------------------------------------------------------------------------------- /compute_modules/client/encoder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import dataclasses 16 | import json 17 | from datetime import date, datetime 18 | from typing import Any 19 | 20 | from compute_modules.function_registry.datetime_conversion_util import DatetimeConversionUtil 21 | 22 | 23 | class CustomJSONEncoder(json.JSONEncoder): 24 | def default(self, obj: Any) -> Any: 25 | if dataclasses.is_dataclass(obj): 26 | return { 27 | field.name: self.default(getattr(obj, field.name)) 28 | for field in dataclasses.fields(obj) 29 | if getattr(obj, field.name) is not None 30 | } 31 | if isinstance(obj, datetime): 32 | return DatetimeConversionUtil.datetime_to_string(obj) 33 | if isinstance(obj, date): 34 | return obj.isoformat() 35 | if isinstance(obj, list): 36 | return [self.default(item) for item in obj] 37 | if isinstance(obj, dict): 38 | return {k: self.default(v) for k, v in obj.items()} 39 | return obj 40 | -------------------------------------------------------------------------------- /tests/infer/test_project/decorated/alias_types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from dataclasses import dataclass 16 | 17 | from compute_modules.annotations import function 18 | from compute_modules.context import QueryContext 19 | from compute_modules.function_registry.types import Byte, Double, Long, Short 20 | 21 | 22 | @dataclass 23 | class LongPayload: 24 | value: Long 25 | 26 | 27 | @function 28 | def return_long(context: QueryContext, event: LongPayload) -> Long: 29 | return event.value 30 | 31 | 32 | @dataclass 33 | class ShortPayload: 34 | value: Short 35 | 36 | 37 | @function 38 | def return_short(context: QueryContext, event: ShortPayload) -> Short: 39 | return event.value 40 | 41 | 42 | @dataclass 43 | class DoublePayload: 44 | value: Double 45 | 46 | 47 | @function 48 | def return_double(context: QueryContext, event: DoublePayload) -> Double: 49 | return event.value 50 | 51 | 52 | @dataclass 53 | class BytePayload: 54 | value: Byte 55 | 56 | 57 | @function 58 | def return_byte(context: QueryContext, event: BytePayload) -> Byte: 59 | return event.value 60 | -------------------------------------------------------------------------------- /compute_modules/function_registry/datetime_conversion_util.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from datetime import datetime, timezone 17 | 18 | 19 | class DatetimeConversionUtil: 20 | DATETIME_FORMAT_STRING = "%Y-%m-%dT%H:%M:%SZ" 21 | DATETIME_FORMAT_HIGHER_PRECISION_STRING = "%Y-%m-%dT%H:%M:%S.%fZ" 22 | 23 | @staticmethod 24 | def datetime_to_string(datetime_obj: datetime) -> str: 25 | obj = datetime_obj.astimezone(timezone.utc) 26 | # Format as ISO 8601 string with 'Z' suffix 27 | return obj.strftime( 28 | DatetimeConversionUtil.DATETIME_FORMAT_HIGHER_PRECISION_STRING 29 | if obj.microsecond > 0 30 | else DatetimeConversionUtil.DATETIME_FORMAT_STRING 31 | ) 32 | 33 | @staticmethod 34 | def string_to_datetime(datetime_string: str) -> datetime: 35 | try: 36 | return datetime.strptime(datetime_string, DatetimeConversionUtil.DATETIME_FORMAT_STRING) 37 | except ValueError: 38 | return datetime.strptime(datetime_string, DatetimeConversionUtil.DATETIME_FORMAT_HIGHER_PRECISION_STRING) 39 | -------------------------------------------------------------------------------- /compute_modules/resources/pipeline_resources.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import json 16 | from os import environ 17 | from typing import Dict 18 | 19 | from .types import PipelineResource 20 | 21 | RESOURCE_ALIAS_MAP = "RESOURCE_ALIAS_MAP" 22 | RESOURCE_ALIAS_NOT_FOUND = """No resource aliases mounted. 23 | This implies the RESOURCE_ALIAS_MAP environment variable has not been set, 24 | or your Compute Module is not running in Pipelines mode. 25 | Please ensure you have set resources mounted on the Compute Module.""" 26 | 27 | 28 | def get_pipeline_resources() -> Dict[str, PipelineResource]: 29 | """Returns a dictionary of resource alias identifier -> Resource. 30 | The identifier(s) in this dict correspond to the identifier used for an input/output 31 | defined in the Configure tab of your 'Pipelines' compute module 32 | """ 33 | if RESOURCE_ALIAS_MAP not in environ: 34 | raise RuntimeError(RESOURCE_ALIAS_NOT_FOUND) 35 | with open(environ[RESOURCE_ALIAS_MAP], encoding="utf-8") as f: 36 | resource_alias_map_raw = json.load(f) 37 | return {key: PipelineResource(**value) for key, value in resource_alias_map_raw.items()} 38 | -------------------------------------------------------------------------------- /compute_modules/function_registry/function.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import functools 16 | from typing import Any, Callable, Dict, List 17 | 18 | from compute_modules.function_registry.function_schema_parser import parse_function_schema 19 | from compute_modules.function_registry.types import ComputeModuleFunctionSchema 20 | 21 | 22 | class Function(Callable[..., Any]): # type: ignore[misc] 23 | def __init__( 24 | self, 25 | function: Callable[..., Any], 26 | edits: List[Any], 27 | ): 28 | self.function = function 29 | self.edits = edits 30 | functools.update_wrapper(self, function) 31 | 32 | def __call__(self, *args, **kwargs): # type: ignore[no-untyped-def] 33 | return self.function(*args, **kwargs) 34 | 35 | def get_function_schema( 36 | self, 37 | api_name_type_id_mapping: Dict[str, str], 38 | ) -> ComputeModuleFunctionSchema: 39 | return parse_function_schema( 40 | function_ref=self.function, 41 | function_name=self.function.__name__, 42 | edits=self.edits, 43 | api_name_type_id_mapping=api_name_type_id_mapping, 44 | throw_on_missing_type_id=True, 45 | ).function_schema 46 | -------------------------------------------------------------------------------- /scripts/ontology/generate_metadata_config.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import json 17 | from typing import Dict, List 18 | 19 | from scripts.ontology._types import ObjectTypeMetadata 20 | from scripts.ontology.metadata_loader import load_object_type_metadata 21 | 22 | 23 | def generate_metadata_config( 24 | foundry_url: str, 25 | token: str, 26 | object_type_rids: List[str], 27 | link_type_rids: List[str], 28 | output_file: str, 29 | ) -> None: 30 | ontology_metadata = load_object_type_metadata( 31 | foundry_url=foundry_url, 32 | token=token, 33 | object_type_rids=object_type_rids, 34 | link_type_rids=link_type_rids, 35 | ) 36 | config = {"apiNameToTypeId": _get_api_name_type_id_mapping(ontology_metadata)} 37 | with open(output_file, "w") as f: 38 | json.dump(config, f) 39 | 40 | 41 | def _get_api_name_type_id_mapping( 42 | ontology_metadata: Dict[str, ObjectTypeMetadata], 43 | ) -> Dict[str, str]: 44 | res = {} 45 | for metadata in ontology_metadata.values(): 46 | if metadata.api_name in res: 47 | raise ValueError(f"Duplicate api name found in ontology metadata: {metadata.api_name}") 48 | res[metadata.api_name] = metadata.type_id 49 | return res 50 | -------------------------------------------------------------------------------- /compute_modules/sources_v2/_back_compat.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import json 17 | import os 18 | from functools import cache 19 | 20 | from ._api import SOURCE_CONFIGURATIONS_PATH, SOURCE_CREDENTIALS_PATH, MountedSourceConfig 21 | 22 | 23 | @cache 24 | def get_mounted_source_secrets() -> dict[str, dict[str, str]]: 25 | creds_path = os.environ.get(SOURCE_CREDENTIALS_PATH) 26 | if not creds_path: 27 | return {} 28 | data = {} 29 | with open(creds_path, "r", encoding="utf-8") as fr: 30 | data.update(json.load(fr)) 31 | if isinstance(data, dict): 32 | return data 33 | else: 34 | raise ValueError("The JSON content is not a dictionary") 35 | 36 | 37 | @cache 38 | def get_mounted_sources() -> dict[str, MountedSourceConfig]: 39 | mounted_sources = {} 40 | configs_path = os.environ.get(SOURCE_CONFIGURATIONS_PATH) 41 | if configs_path: 42 | with open(configs_path, "r", encoding="utf-8") as fr: 43 | raw_configs = json.load(fr) 44 | if isinstance(raw_configs, dict): 45 | mounted_sources = {key: MountedSourceConfig.from_dict(value) for key, value in raw_configs.items()} 46 | else: 47 | raise ValueError("The JSON content is not a dictionary") 48 | return mounted_sources 49 | -------------------------------------------------------------------------------- /tests/infer/test_project/manually_registered/mixed_registration.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from dataclasses import dataclass 16 | from typing import Set, TypedDict 17 | 18 | from compute_modules.context import QueryContext 19 | from compute_modules.function_registry.function_registry import add_function, add_functions 20 | from compute_modules.function_registry.types import Byte 21 | from compute_modules.startup import start_compute_module 22 | 23 | 24 | @dataclass 25 | class Mixed1Wrapper: 26 | value: int 27 | 28 | 29 | @dataclass 30 | class Mixed2Wrapper: 31 | value: Set[int] 32 | 33 | 34 | @dataclass 35 | class Mixed3Wrapper: 36 | value: Byte 37 | 38 | 39 | class Mixed4Wrapper(TypedDict): 40 | name: str 41 | 42 | 43 | def mixed_1(context: QueryContext, wrapper: Mixed1Wrapper) -> int: 44 | return wrapper.value 45 | 46 | 47 | def mixed_2(context: QueryContext, wrapper: Mixed2Wrapper) -> Set[int]: 48 | return wrapper.value 49 | 50 | 51 | def mixed_3(context: QueryContext, wrapper: Mixed3Wrapper) -> Byte: 52 | return wrapper.value 53 | 54 | 55 | def mixed_4(context: QueryContext, wrapper: Mixed4Wrapper) -> str: 56 | return wrapper["name"] 57 | 58 | 59 | if __name__ == "__main__": 60 | add_function(mixed_1) 61 | add_functions(mixed_2, mixed_3, mixed_4) 62 | start_compute_module() 63 | -------------------------------------------------------------------------------- /scripts/ontology/metadata_client.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from typing import List, Tuple 17 | 18 | import requests 19 | 20 | from scripts.ontology._types import OntologyMetadataLinkTypeOuter, OntologyMetadataObjectTypeOuter 21 | 22 | 23 | class OntologyMetadataClient: 24 | bulk_load_entities_path = "{foundry_url}/ontology-metadata/api/ontology/ontology/bulkLoadEntities" 25 | 26 | def __init__(self, foundry_url: str, token: str): 27 | self.foundry_url = foundry_url 28 | self.token = token 29 | 30 | def bulk_load_entities( 31 | self, 32 | object_type_rids: List[str], 33 | link_type_rids: List[str], 34 | ) -> Tuple[List[OntologyMetadataObjectTypeOuter], List[OntologyMetadataLinkTypeOuter]]: 35 | payload = { 36 | "objectTypes": [ 37 | {"identifier": {"objectTypeRid": rid, "type": "objectTypeRid"}} for rid in object_type_rids 38 | ], 39 | "linkTypes": [{"identifier": {"linkTypeRid": rid, "type": "linkTypeRid"}} for rid in link_type_rids], 40 | "loadRedacted": True, 41 | "includeObjectTypesWithoutSearchableDatasources": True, 42 | "includeEntityMetadata": False, 43 | } 44 | response = requests.post( 45 | self.bulk_load_entities_path.format(foundry_url=self.foundry_url), 46 | headers={ 47 | "Authorization": f"Bearer {self.token}", 48 | }, 49 | json=payload, 50 | ) 51 | result = response.json() 52 | return result["objectTypes"], result["linkTypes"] 53 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "foundry-compute-modules" 3 | version = "0.0.0" 4 | description = "The official Python library for creating Compute Modules" 5 | authors = ["Palantir Technologies, Inc."] 6 | license = "Apache-2.0" 7 | readme = "README.md" 8 | repository = "https://github.com/palantir/python-compute-module" 9 | keywords = ["Palantir", "Foundry", "Compute Modules"] 10 | packages = [{ include = "compute_modules" }] 11 | 12 | [tool.poetry.dependencies] 13 | python = "^3.9" 14 | requests = "^2.32.3" 15 | external-systems = { version = "^0.107.0", optional = true } 16 | pyyaml = { version = "^6.0.1", optional = true } 17 | 18 | 19 | [tool.poetry.extras] 20 | sources = ["external-systems", "pyyaml"] 21 | 22 | 23 | [tool.poetry.group.dev.dependencies] 24 | black = "24.1.1" 25 | isort = "5.13.2" 26 | mypy = "1.9.0" 27 | ruff = "0.9.7" 28 | pytest = "8.0.0" 29 | pytest-html = "4.1.1" 30 | poethepoet = "^0.29.0" 31 | types-requests = "^2.32.0.20241016" 32 | 33 | [build-system] 34 | requires = ["poetry-core"] 35 | build-backend = "poetry.core.masonry.api" 36 | 37 | [tool.poe.tasks] 38 | test = "python -c 'from scripts.checks import test; test()'" 39 | check_format = "python -c 'from scripts.checks import check_format; check_format()'" 40 | check_mypy = "python -c 'from scripts.checks import check_mypy; check_mypy()'" 41 | format = "python -c 'from scripts.checks import format; format()'" 42 | check_license = "python -c 'from scripts.checks import check_license; check_license()'" 43 | license = "python -c 'from scripts.checks import license; license()'" 44 | set_version = "python -c 'from scripts.set_version import main; main()'" 45 | 46 | [tool.poetry.scripts] 47 | cm_ontology_metadata_config = "compute_modules.bin.ontology.generate_metadata_config:main" 48 | cm_python_function_infer = "compute_modules.bin.static_inference.infer:main" 49 | 50 | [tool.black] 51 | line_length = 120 52 | 53 | [tool.isort] 54 | profile = "black" 55 | line_length = 120 56 | 57 | [tool.pytest.ini_options] 58 | addopts = "--junitxml=./build/pytest-results/pytest-results.xml --html=./build/pytest-results/pytest-results.html" 59 | cache_dir = "build/.pytest_cache" 60 | testpaths = ["tests"] 61 | 62 | [tool.ruff] 63 | line-length = 120 64 | cache-dir = "build/.ruff_cache" 65 | 66 | [tool.mypy] 67 | strict = true 68 | -------------------------------------------------------------------------------- /scripts/ontology/main.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import argparse 17 | 18 | from scripts.ontology._config_path import get_ontology_config_file 19 | from scripts.ontology.generate_metadata_config import generate_metadata_config 20 | 21 | 22 | def main() -> None: 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument("source") 25 | parser.add_argument( 26 | "--object-type-rid", 27 | required=False, 28 | nargs="*", 29 | dest="object_type_rids", 30 | default=[], 31 | ) 32 | parser.add_argument( 33 | "--link-type-rid", 34 | required=False, 35 | nargs="*", 36 | dest="link_type_rids", 37 | default=[], 38 | ) 39 | parser.add_argument( 40 | "-t", 41 | "--token", 42 | required=True, 43 | default=None, 44 | ) 45 | parser.add_argument( 46 | "--foundry-url", 47 | required=True, 48 | default=None, 49 | ) 50 | parser.add_argument( 51 | "--ontology-metadata-config", 52 | required=False, 53 | dest="ontology_metadata_config_file", 54 | default=None, 55 | ) 56 | arguments = parser.parse_args() 57 | output_file = get_ontology_config_file(arguments.ontology_metadata_config_file) 58 | print("Generating config file...") 59 | generate_metadata_config( 60 | foundry_url=arguments.foundry_url, 61 | token=arguments.token, 62 | object_type_rids=arguments.object_type_rids, 63 | link_type_rids=arguments.link_type_rids, 64 | output_file=output_file, 65 | ) 66 | print(f"Wrote config file to {output_file}") 67 | 68 | 69 | if __name__ == "__main__": 70 | main() 71 | -------------------------------------------------------------------------------- /compute_modules/context/types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from dataclasses import dataclass 17 | from typing import Any, Dict, Optional 18 | 19 | 20 | @dataclass 21 | class QueryContext: 22 | """Metadata for the job being executed that is not included in the event payload""" 23 | 24 | authHeader: str 25 | """Foundry auth token that can be used to call other services within Foundry. 26 | Only available in certain modes 27 | """ 28 | 29 | jobId: str 30 | """The unique identifier for the given job""" 31 | 32 | tempCredsAuthToken: Optional[str] = None 33 | """A temporary token that is used with the Foundry data sidecar.""" 34 | 35 | CLIENT_ID: Optional[str] = None 36 | """Client ID of the third party application associated with this compute module. 37 | Present if compute module is configured to have application's permissions. 38 | Use this to get a Foundry scoped token from your third party application service user. 39 | """ 40 | 41 | CLIENT_SECRET: Optional[str] = None 42 | """Client secret of the third party application associated with this compute module. 43 | Present if compute module is configured to have application's permissions. 44 | Use this to get a Foundry scoped token from your third party application service user. 45 | """ 46 | 47 | sources: Optional[Dict[str, Any]] = None 48 | """dict containing the secrets of any sources configured for this compute module.""" 49 | 50 | source_configs: Optional[Dict[str, Any]] = None 51 | """dict containing the configuration of any sources configured for this compute module.""" 52 | 53 | userId: Optional[str] = None 54 | """The unique identifier for the user who initiated the job""" 55 | -------------------------------------------------------------------------------- /scripts/ontology/_types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from dataclasses import dataclass 17 | from typing import Dict, List, Optional, TypedDict 18 | 19 | # Not all fields are present here because I manually defined these and am lazy; 20 | # Ideally we'd be able to pull in conjure types here but I'm not aware of how to do so while 21 | # staying within the bounds of what is allowed of OSS. Open to suggestions! 22 | 23 | 24 | @dataclass(frozen=True) 25 | class ObjectTypeMetadata: 26 | api_name: str 27 | type_id: str 28 | primary_key_id: str 29 | property_api_name_to_id: Dict[str, str] 30 | link_type_api_name_to_id: Dict[str, str] 31 | 32 | 33 | # Object types 34 | class ObjectyPropertyType(TypedDict): 35 | id: str 36 | rid: str 37 | apiName: str 38 | 39 | 40 | class OntologyMetadataObjectType(TypedDict): 41 | rid: str 42 | apiName: str 43 | id: str 44 | propertyTypes: Dict[str, ObjectyPropertyType] 45 | primaryKeys: List[str] 46 | 47 | 48 | class OntologyMetadataObjectTypeOuter(TypedDict): 49 | objectType: OntologyMetadataObjectType 50 | 51 | 52 | # Link types 53 | class ManyToManyMetadata(TypedDict): 54 | apiName: str 55 | 56 | 57 | class LinkTypeManyToMany(TypedDict): 58 | objectTypeRidA: str 59 | objectTypeRidB: str 60 | objectTypeAToBLinkMetadata: ManyToManyMetadata 61 | objectTypeBToALinkMetadata: ManyToManyMetadata 62 | 63 | 64 | class OntologyLinkTypeDefinition(TypedDict): 65 | type: str 66 | manyToMany: Optional[LinkTypeManyToMany] 67 | 68 | 69 | class OntologyMetadataLinkType(TypedDict): 70 | id: str 71 | definition: OntologyLinkTypeDefinition 72 | 73 | 74 | class OntologyMetadataLinkTypeOuter(TypedDict): 75 | linkType: OntologyMetadataLinkType 76 | -------------------------------------------------------------------------------- /tests/infer/test_project/manually_registered/multiple_functions.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import datetime 16 | from dataclasses import dataclass 17 | from typing import List 18 | 19 | from compute_modules.context import QueryContext 20 | from compute_modules.function_registry.function_registry import add_function, add_functions 21 | from compute_modules.startup import start_compute_module 22 | from tests.infer.test_project._types import DummyOntologyType, OntologyEdit 23 | 24 | 25 | @dataclass 26 | class Point: 27 | x: float 28 | y: float 29 | z: float 30 | 31 | 32 | @dataclass 33 | class Message: 34 | message: str 35 | from_id: int 36 | to_id: int 37 | timestamp: datetime.datetime 38 | 39 | 40 | @dataclass 41 | class Messages: 42 | messages: List[Message] 43 | 44 | 45 | @dataclass 46 | class MultipleEventWrapper: 47 | points: List[Point] 48 | messages: Messages 49 | 50 | 51 | @dataclass 52 | class Number: 53 | value: int 54 | 55 | 56 | def return_complex_in_main(context: QueryContext, event: MultipleEventWrapper) -> str: 57 | return f"Points: {event.points}, Messages: {event.messages}" 58 | 59 | 60 | def return_number_in_main(context: QueryContext, number: Number) -> Number: 61 | return number.value # type: ignore[return-value] 62 | 63 | 64 | @dataclass 65 | class OntologyEditParams: 66 | name: str 67 | 68 | 69 | def ontology_edit_function( 70 | context: QueryContext, 71 | event: OntologyEditParams, 72 | ) -> list[OntologyEdit]: 73 | return [] 74 | 75 | 76 | EDIT_TYPES = [DummyOntologyType] 77 | 78 | 79 | if __name__ == "__main__": 80 | add_functions(return_complex_in_main, return_number_in_main) 81 | add_function(ontology_edit_function, edits=EDIT_TYPES) 82 | start_compute_module() 83 | -------------------------------------------------------------------------------- /.circleci/config.yml: -------------------------------------------------------------------------------- 1 | palantir_aliases: 2 | - &always-run 3 | filters: 4 | branches: 5 | only: /.*/ 6 | tags: 7 | only: /.*/ 8 | 9 | version: 2.1 10 | jobs: 11 | test: 12 | parameters: 13 | python_version: 14 | type: string 15 | docker: 16 | - image: cimg/python:<< parameters.python_version >> 17 | steps: 18 | - checkout 19 | - run: poetry install --extras sources 20 | - run: poetry run poe test 21 | 22 | mypy: 23 | docker: 24 | - image: cimg/python:3.12 25 | steps: 26 | - checkout 27 | - run: poetry install --extras sources 28 | - run: poetry run poe check_mypy 29 | 30 | check_format: 31 | docker: 32 | - image: cimg/python:3.12 33 | steps: 34 | - checkout 35 | - run: poetry install --extras sources 36 | - run: poetry run poe check_format 37 | 38 | check_license: 39 | docker: 40 | - image: cimg/python:3.12 41 | steps: 42 | - checkout 43 | - run: poetry install --extras sources 44 | - run: poetry run poe check_license 45 | 46 | circle-all: 47 | docker: 48 | - image: node:lts 49 | steps: 50 | - run: echo "Done!" 51 | 52 | publish: 53 | docker: 54 | - image: cimg/python:3.12 55 | steps: 56 | - checkout 57 | - run: poetry install --extras sources 58 | - run: poetry run poe set_version 59 | - run: poetry publish -v -u $PYPI_USERNAME -p $PYPI_PASSWORD --build 60 | 61 | workflows: 62 | version: 2 63 | build: 64 | jobs: 65 | - test: 66 | <<: *always-run 67 | name: python-<< matrix.python_version>> 68 | matrix: 69 | parameters: 70 | python_version: ["3.9", "3.10", "3.11", "3.12"] 71 | - mypy: 72 | <<: *always-run 73 | - check_format: 74 | <<: *always-run 75 | - check_license: 76 | <<: *always-run 77 | - circle-all: 78 | <<: *always-run 79 | requires: 80 | - python-3.9 81 | - python-3.10 82 | - python-3.11 83 | - python-3.12 84 | - mypy 85 | - check_format 86 | - check_license 87 | - publish: 88 | requires: 89 | - circle-all 90 | filters: 91 | tags: { only: '/^[0-9]+(\.[0-9]+)+(-[a-zA-Z]+[0-9]*)*$/' } 92 | branches: { ignore: /.*/ } -------------------------------------------------------------------------------- /compute_modules/bin/ontology/_types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from dataclasses import dataclass 17 | from typing import Dict, List, Optional, TypedDict 18 | 19 | # Not all fields are present here because I manually defined these and am lazy; 20 | # Ideally we'd be able to pull in conjure types here but I'm not aware of how to do so while 21 | # staying within the bounds of what is allowed of OSS. Open to suggestions! 22 | 23 | 24 | @dataclass(frozen=True) 25 | class ObjectTypeMetadata: 26 | objectTypeApiName: str 27 | objectTypeId: str 28 | primaryKeyPropertyId: str 29 | properties: Dict[str, str] 30 | links: Dict[str, str] 31 | 32 | 33 | @dataclass(frozen=True) 34 | class RuntimeMetadata: 35 | objectMetadata: Dict[str, ObjectTypeMetadata] 36 | 37 | 38 | # Object types 39 | class ObjectyPropertyType(TypedDict): 40 | id: str 41 | rid: str 42 | apiName: str 43 | 44 | 45 | class OntologyMetadataObjectType(TypedDict): 46 | rid: str 47 | apiName: str 48 | id: str 49 | propertyTypes: Dict[str, ObjectyPropertyType] 50 | primaryKeys: List[str] 51 | 52 | 53 | class OntologyMetadataObjectTypeOuter(TypedDict): 54 | objectType: OntologyMetadataObjectType 55 | 56 | 57 | # Link types 58 | class ManyToManyMetadata(TypedDict): 59 | apiName: str 60 | 61 | 62 | class LinkTypeManyToMany(TypedDict): 63 | objectTypeRidA: str 64 | objectTypeRidB: str 65 | objectTypeAToBLinkMetadata: ManyToManyMetadata 66 | objectTypeBToALinkMetadata: ManyToManyMetadata 67 | 68 | 69 | class OntologyLinkTypeDefinition(TypedDict): 70 | type: str 71 | manyToMany: Optional[LinkTypeManyToMany] 72 | 73 | 74 | class OntologyMetadataLinkType(TypedDict): 75 | id: str 76 | definition: OntologyLinkTypeDefinition 77 | 78 | 79 | class OntologyMetadataLinkTypeOuter(TypedDict): 80 | linkType: OntologyMetadataLinkType 81 | -------------------------------------------------------------------------------- /compute_modules/bin/ontology/metadata_client.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from typing import List, Tuple 17 | from urllib.parse import urlparse 18 | 19 | import requests 20 | 21 | from compute_modules.bin.ontology._types import OntologyMetadataLinkTypeOuter, OntologyMetadataObjectTypeOuter 22 | 23 | 24 | def _clean_url(url: str) -> str: 25 | parsed_url = urlparse(url) 26 | if not parsed_url.scheme: 27 | parsed_url = urlparse(f"https://{url}") 28 | return f"https://{parsed_url.netloc}" 29 | 30 | 31 | class OntologyMetadataClient: 32 | bulk_load_entities_path = "{foundry_url}/ontology-metadata/api/ontology/ontology/bulkLoadEntities" 33 | 34 | def __init__(self, foundry_url: str, token: str): 35 | self.foundry_url = _clean_url(foundry_url) 36 | self.token = token 37 | 38 | def bulk_load_entities( 39 | self, 40 | object_type_rids: List[str], 41 | link_type_rids: List[str], 42 | ) -> Tuple[List[OntologyMetadataObjectTypeOuter], List[OntologyMetadataLinkTypeOuter]]: 43 | payload = { 44 | "objectTypes": [ 45 | {"identifier": {"objectTypeRid": rid, "type": "objectTypeRid"}} for rid in object_type_rids 46 | ], 47 | "linkTypes": [{"identifier": {"linkTypeRid": rid, "type": "linkTypeRid"}} for rid in link_type_rids], 48 | "loadRedacted": True, 49 | "includeObjectTypesWithoutSearchableDatasources": True, 50 | "includeEntityMetadata": False, 51 | } 52 | response = requests.post( 53 | self.bulk_load_entities_path.format(foundry_url=self.foundry_url), 54 | headers={ 55 | "Authorization": f"Bearer {self.token}", 56 | }, 57 | json=payload, 58 | ) 59 | response.raise_for_status() 60 | result = response.json() 61 | return result["objectTypes"], result["linkTypes"] 62 | -------------------------------------------------------------------------------- /tests/function_registry/dummy_app_with_issues.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from typing import TypedDict 17 | 18 | 19 | # Failure case 1 20 | class BadClassNoTypeHints: 21 | def __init__(self, arg1: str, arg2: int): 22 | self.arg1 = arg1 23 | self.arg2 = arg2 24 | 25 | 26 | class NoTypeHintsWrapper(TypedDict): 27 | param: BadClassNoTypeHints 28 | 29 | 30 | def dummy_no_type_hints(context, event: NoTypeHintsWrapper) -> int: # type: ignore[no-untyped-def] 31 | return 1 32 | 33 | 34 | # Failure case 2 35 | class BadClassNoInitHints: 36 | arg1: str 37 | arg2: int 38 | 39 | def __init__(self, arg1, arg2) -> None: # type: ignore[no-untyped-def] 40 | self.arg1 = arg1 41 | self.arg2 = arg2 42 | 43 | 44 | class NoInitHintsWrapper(TypedDict): 45 | param: BadClassNoInitHints 46 | 47 | 48 | def dummy_no_init_hints(context, event: NoInitHintsWrapper) -> int: # type: ignore[no-untyped-def] 49 | return 1 50 | 51 | 52 | # Failure case 3 53 | class BadClassArgsInit: 54 | arg1: str 55 | arg2: int 56 | 57 | def __init__(self, arg1: str, arg2: int, *args) -> None: # type: ignore[no-untyped-def] 58 | self.arg1 = arg1 59 | self.arg2 = arg2 60 | 61 | 62 | class ArgsInitWrapper(TypedDict): 63 | param: BadClassArgsInit 64 | 65 | 66 | def dummy_args_init(context, event: ArgsInitWrapper) -> int: # type: ignore[no-untyped-def] 67 | return 1 68 | 69 | 70 | # Failure case 4 71 | class BadClassKwargsInit: 72 | arg1: str 73 | arg2: int 74 | 75 | def __init__(self, arg1: str, arg2: int, **kwargs) -> None: # type: ignore[no-untyped-def] 76 | self.arg1 = kwargs["arg1"] 77 | self.arg2 = kwargs["arg2"] 78 | 79 | 80 | class KwargsInitWrapper(TypedDict): 81 | param: BadClassKwargsInit 82 | 83 | 84 | def dummy_kwargs_init(context, event: KwargsInitWrapper) -> int: # type: ignore[no-untyped-def] 85 | return 1 86 | -------------------------------------------------------------------------------- /tests/infer/test_project/decorated/primitive_types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import datetime 16 | import decimal 17 | from dataclasses import dataclass 18 | 19 | from compute_modules.annotations import function 20 | from compute_modules.context import QueryContext 21 | 22 | 23 | @dataclass 24 | class BytesWrapper: 25 | value: bytes 26 | 27 | 28 | @function 29 | def return_bytes(context: QueryContext, wrapper: BytesWrapper) -> bytes: 30 | return wrapper.value 31 | 32 | 33 | @dataclass 34 | class BoolWrapper: 35 | value: bool 36 | 37 | 38 | @function 39 | def return_bool(context: QueryContext, wrapper: BoolWrapper) -> bool: 40 | return wrapper.value 41 | 42 | 43 | @dataclass 44 | class DateWrapper: 45 | value: datetime.date 46 | 47 | 48 | @function 49 | def return_date(context: QueryContext, wrapper: DateWrapper) -> datetime.date: 50 | return wrapper.value 51 | 52 | 53 | @dataclass 54 | class DecimalWrapper: 55 | value: decimal.Decimal 56 | 57 | 58 | @function 59 | def return_decimal(context: QueryContext, wrapper: DecimalWrapper) -> decimal.Decimal: 60 | return wrapper.value 61 | 62 | 63 | @dataclass 64 | class FloatWrapper: 65 | value: float 66 | 67 | 68 | @function 69 | def return_float(context: QueryContext, wrapper: FloatWrapper) -> float: 70 | return wrapper.value 71 | 72 | 73 | @dataclass 74 | class IntWrapper: 75 | value: int 76 | 77 | 78 | @function 79 | def return_int(context: QueryContext, wrapper: IntWrapper) -> int: 80 | return wrapper.value 81 | 82 | 83 | @dataclass 84 | class StrWrapper: 85 | value: str 86 | 87 | 88 | @function 89 | def return_str(context: QueryContext, wrapper: StrWrapper) -> str: 90 | return wrapper.value 91 | 92 | 93 | @dataclass 94 | class DatetimeWrapper: 95 | value: datetime.datetime 96 | 97 | 98 | @function 99 | def return_datetime(context: QueryContext, wrapper: DatetimeWrapper) -> datetime.datetime: 100 | return wrapper.value 101 | -------------------------------------------------------------------------------- /compute_modules/function_registry/types.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import datetime 17 | import decimal 18 | import typing 19 | from dataclasses import dataclass 20 | 21 | DataTypeDict = typing.Dict[str, typing.Any] 22 | 23 | 24 | class DataType(typing.TypedDict): 25 | """Data type schema""" 26 | 27 | dataType: DataTypeDict 28 | 29 | 30 | class PythonClassNode(typing.TypedDict): 31 | """Represents a Python class constructor with a dict containing refs to any child fields that are also custom classes""" 32 | 33 | constructor: typing.Any 34 | children: typing.Optional[typing.Dict[str, "PythonClassNode"]] 35 | 36 | 37 | class FunctionInputType(DataType): 38 | """Function input schema""" 39 | 40 | name: str 41 | required: typing.Literal[True] 42 | constraints: typing.List[typing.Any] 43 | 44 | 45 | class FunctionOutputType(typing.TypedDict): 46 | """Function output schema""" 47 | 48 | type: typing.Literal["single"] 49 | single: DataType 50 | 51 | 52 | class FunctionOntologyProvenance(typing.TypedDict): 53 | editedObjects: typing.Dict[str, typing.Dict[None, None]] 54 | editedLinks: typing.Dict[str, typing.Dict[None, None]] 55 | 56 | 57 | class ComputeModuleFunctionSchema(typing.TypedDict): 58 | """Represents the function schema for a Compute Module function""" 59 | 60 | # TODO: implement apiName & namespace 61 | functionName: str 62 | inputs: typing.List[FunctionInputType] 63 | output: FunctionOutputType 64 | ontologyProvenance: typing.Optional[FunctionOntologyProvenance] 65 | 66 | 67 | @dataclass 68 | class ParseFunctionSchemaResult: 69 | function_schema: ComputeModuleFunctionSchema 70 | class_node: typing.Optional[PythonClassNode] 71 | is_context_typed: bool 72 | 73 | 74 | Byte = typing.NewType("Byte", int) 75 | Double = typing.NewType("Double", float) 76 | Long = typing.NewType("Long", int) 77 | Short = typing.NewType("Short", int) 78 | 79 | AllowedKeyTypes = typing.Union[ 80 | bytes, 81 | bool, 82 | Byte, 83 | datetime.date, 84 | decimal.Decimal, 85 | Double, 86 | float, 87 | int, 88 | Long, 89 | # TODO: handle ontology types? 90 | # OntologyObject 91 | Short, 92 | str, 93 | datetime.datetime, 94 | ] 95 | -------------------------------------------------------------------------------- /tests/auth/test_third_party.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from typing import Generator, List 17 | from unittest.mock import MagicMock, patch 18 | 19 | import pytest 20 | 21 | from compute_modules.auth import RefreshingOauthToken 22 | 23 | # Mock data 24 | MOCK_HOSTNAME: str = "example.com" 25 | MOCK_SCOPE: List[str] = ["api:ontologies-read"] 26 | MOCK_TOKEN: str = "mock_token" 27 | 28 | 29 | @pytest.fixture 30 | def mock_oauth() -> Generator[MagicMock, None, None]: 31 | with patch("compute_modules.auth.third_party.oauth") as mock_oauth: 32 | yield mock_oauth 33 | 34 | 35 | @pytest.fixture 36 | def mock_time() -> Generator[MagicMock, None, None]: 37 | with patch("time.time") as mock_time: 38 | yield mock_time 39 | 40 | 41 | def test_get_token_initial_fetch(mock_oauth: MagicMock, mock_time: MagicMock) -> None: 42 | mock_oauth.return_value = MOCK_TOKEN 43 | mock_time.return_value = 1000 44 | 45 | token_refresher = RefreshingOauthToken(MOCK_HOSTNAME, MOCK_SCOPE) 46 | token: str = token_refresher.get_token() 47 | 48 | assert token == MOCK_TOKEN 49 | mock_oauth.assert_called_once_with(MOCK_HOSTNAME, MOCK_SCOPE) 50 | 51 | 52 | def test_get_token_refresh_needed(mock_oauth: MagicMock, mock_time: MagicMock) -> None: 53 | mock_oauth.return_value = MOCK_TOKEN 54 | initial_time: int = 1000 55 | mock_time.side_effect = [initial_time, initial_time + 1900] # Outside refresh interval 56 | 57 | token_refresher = RefreshingOauthToken(MOCK_HOSTNAME, MOCK_SCOPE) 58 | token_refresher.get_token() 59 | token: str = token_refresher.get_token() # Should trigger refresh 60 | 61 | assert token == MOCK_TOKEN 62 | assert mock_oauth.call_count == 2 63 | 64 | 65 | def test_get_token_no_refresh_needed(mock_oauth: MagicMock, mock_time: MagicMock) -> None: 66 | mock_oauth.return_value = MOCK_TOKEN 67 | initial_time: int = 1000 68 | mock_time.side_effect = [initial_time, initial_time + 1700] # Within refresh interval 69 | 70 | token_refresher = RefreshingOauthToken(MOCK_HOSTNAME, MOCK_SCOPE) 71 | token_refresher.get_token() 72 | token: str = token_refresher.get_token() # Should not trigger refresh 73 | 74 | assert token == MOCK_TOKEN 75 | assert mock_oauth.call_count == 1 76 | -------------------------------------------------------------------------------- /compute_modules/function_registry/function_payload_converter.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import datetime 17 | import logging 18 | import typing 19 | 20 | from .types import PythonClassNode 21 | 22 | logger = logging.getLogger(__name__) 23 | 24 | 25 | def convert_payload( 26 | raw_payload: typing.Any, 27 | class_tree: PythonClassNode, 28 | ) -> typing.Any: 29 | try: 30 | if raw_payload is None: 31 | return None 32 | # No children indicates raw_payload should be a primtive type 33 | if class_tree["children"] is None: 34 | return class_tree["constructor"](raw_payload) 35 | type_constructor = class_tree["constructor"] 36 | if type_constructor is datetime.datetime: 37 | return datetime.datetime.fromisoformat(raw_payload.replace("Z", "+00:00")) 38 | if type_constructor is list: 39 | child_class_tree = class_tree["children"]["list"] 40 | return list([convert_payload(el, child_class_tree) for el in raw_payload]) 41 | if type_constructor is dict: 42 | key_class_tree = class_tree["children"]["key"] 43 | value_class_tree = class_tree["children"]["value"] 44 | return { 45 | convert_payload(key, key_class_tree): convert_payload(value, value_class_tree) 46 | for key, value in raw_payload.items() 47 | } 48 | if type_constructor is typing.Optional: 49 | return raw_payload 50 | if type_constructor is set: 51 | child_class_tree = class_tree["children"]["set"] 52 | return set([convert_payload(el, child_class_tree) for el in raw_payload]) 53 | # Complex class 54 | converted_children = {} 55 | for child_key, child_class_tree in class_tree["children"].items(): 56 | # if child is optional and no value provided, default to None 57 | if child_class_tree["constructor"] is typing.Optional and child_key not in raw_payload: 58 | raw_payload[child_key] = None 59 | converted_children[child_key] = convert_payload(raw_payload[child_key], child_class_tree) 60 | return type_constructor(**converted_children) 61 | except Exception as e: 62 | logger.error(f"Error converting {raw_payload} to type {class_tree['constructor']}") 63 | raise e 64 | -------------------------------------------------------------------------------- /tests/function_registry/dummy_app.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from dataclasses import dataclass 17 | from datetime import date, datetime 18 | from decimal import Decimal 19 | from typing import Dict, Iterable, List, Optional, Set, Union 20 | 21 | from compute_modules.context import QueryContext 22 | 23 | 24 | @dataclass 25 | class ChildClass: 26 | timestamp: datetime 27 | some_value: float 28 | another_optional_field: Union[str, None] 29 | 30 | 31 | class ParentClass: 32 | some_flag: bool 33 | some_value: int 34 | child: ChildClass 35 | 36 | def __init__(self, some_flag: bool, some_value: int, child: ChildClass) -> None: 37 | self.some_flag = some_flag 38 | self.some_value = some_value 39 | self.child = child 40 | 41 | 42 | @dataclass 43 | class DummyInput: 44 | parent_class: ParentClass 45 | optional_field: Optional[str] 46 | set_field: Set[date] 47 | map_field: Dict[bytes, Decimal] 48 | datetime_list: List[datetime] 49 | some_flag: bool 50 | optional_default_value_field: Optional[str] = None 51 | 52 | 53 | @dataclass 54 | class DummyOutput: 55 | res1: bool 56 | res2: Dict[str, float] 57 | res3: List[datetime] 58 | 59 | 60 | @dataclass 61 | class ClassWithBareDict: 62 | dict_field: dict # type: ignore[type-arg] 63 | 64 | 65 | def dummy_func_1(context, event: DummyInput) -> DummyOutput: # type: ignore[no-untyped-def] 66 | """Example function with type hints""" 67 | key_value = event.optional_field or "default" 68 | return DummyOutput(res1=True, res2={key_value: event.parent_class.child.some_value}, res3=event.datetime_list) 69 | 70 | 71 | def dummy_func_2(context, event): # type: ignore[no-untyped-def] 72 | """No type hints example""" 73 | return event["blah"] 74 | 75 | 76 | def dummy_func_3(context: QueryContext, event) -> int: # type: ignore[no-untyped-def] 77 | """Example function with type hint for context & return type only""" 78 | return 1 79 | 80 | 81 | def dummy_func_4(context: QueryContext, event: ClassWithBareDict) -> int: 82 | """Example function with type hint for context & return type only""" 83 | return 1 84 | 85 | 86 | def dummy_func_5(context, event) -> Iterable[str]: # type: ignore[no-untyped-def] 87 | """Example function with type hint for generator return type""" 88 | for i in range(10): 89 | yield f"string {i}" 90 | -------------------------------------------------------------------------------- /compute_modules/function_registry/function_registry.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from typing import Any, Callable, Dict, List, Optional 17 | 18 | from .function_schema_parser import parse_function_schema 19 | from .types import ComputeModuleFunctionSchema, PythonClassNode 20 | 21 | REGISTERED_FUNCTIONS: Dict[str, Callable[..., Any]] = {} 22 | FUNCTION_SCHEMAS: List[ComputeModuleFunctionSchema] = [] 23 | FUNCTION_SCHEMA_CONVERSIONS: Dict[str, PythonClassNode] = {} 24 | IS_FUNCTION_CONTEXT_TYPED: Dict[str, bool] = {} 25 | STREAMING: Dict[str, bool] = {} 26 | 27 | 28 | def add_functions(*args: Callable[..., Any]) -> None: 29 | for function_ref in args: 30 | add_function(function_ref=function_ref) 31 | 32 | 33 | def add_function( 34 | function_ref: Callable[..., Any], 35 | *, 36 | streaming: bool = False, 37 | edits: Optional[List[Any]] = None, 38 | ) -> None: 39 | """Parse & register a Compute Module function""" 40 | function_name = function_ref.__name__ 41 | parse_result = parse_function_schema( 42 | function_ref, 43 | function_name, 44 | edits=edits if edits is not None else [], 45 | # TODO: currently do not support runtime function inference for OntologyEdits. 46 | # Not sure if we will but would need to update here if we decide to do so 47 | api_name_type_id_mapping={}, 48 | ) 49 | _register_parsed_function( 50 | function_name=function_name, 51 | function_ref=function_ref, 52 | function_schema=parse_result.function_schema, 53 | function_schema_conversion=parse_result.class_node, 54 | is_context_typed=parse_result.is_context_typed, 55 | streaming=streaming, 56 | ) 57 | 58 | 59 | def _register_parsed_function( 60 | function_name: str, 61 | function_ref: Callable[..., Any], 62 | function_schema: ComputeModuleFunctionSchema, 63 | function_schema_conversion: Optional[PythonClassNode], 64 | is_context_typed: bool, 65 | streaming: bool, 66 | ) -> None: 67 | """Registers a Compute Module function""" 68 | REGISTERED_FUNCTIONS[function_name] = function_ref 69 | FUNCTION_SCHEMAS.append(function_schema) 70 | IS_FUNCTION_CONTEXT_TYPED[function_name] = is_context_typed 71 | STREAMING[function_name] = streaming 72 | if function_schema_conversion is not None: 73 | FUNCTION_SCHEMA_CONVERSIONS[function_name] = function_schema_conversion 74 | -------------------------------------------------------------------------------- /compute_modules/auth/third_party.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import http.client 17 | import json 18 | import os 19 | import ssl 20 | import time 21 | import urllib.parse 22 | from typing import Any, List, Optional, Tuple 23 | 24 | 25 | def retrieve_third_party_id_and_creds() -> Tuple[Optional[str], Optional[str]]: 26 | CLIENT_ID = os.getenv("CLIENT_ID") 27 | CLIENT_SECRET = os.getenv("CLIENT_SECRET") 28 | return CLIENT_ID, CLIENT_SECRET 29 | 30 | 31 | def oauth(hostname: str, scope: List[str]) -> Any: 32 | CLIENT_ID, CLIENT_SECRET = retrieve_third_party_id_and_creds() 33 | if CLIENT_ID and CLIENT_SECRET: 34 | params = urllib.parse.urlencode( 35 | { 36 | "grant_type": "client_credentials", 37 | "client_id": CLIENT_ID, 38 | "client_secret": CLIENT_SECRET, 39 | "scope": " ".join(scope), 40 | } 41 | ) 42 | headers = { 43 | "Content-Type": "application/x-www-form-urlencoded", 44 | } 45 | 46 | ssl_context = ssl.create_default_context(cafile=os.environ.get("DEFAULT_CA_PATH")) 47 | conn = http.client.HTTPSConnection(hostname, context=ssl_context) 48 | conn.request("POST", "/multipass/api/oauth2/token", params, headers) 49 | response = conn.getresponse() 50 | data = response.read() 51 | conn.close() 52 | if response.status == 200: 53 | try: 54 | token_data = json.loads(data) 55 | if isinstance(token_data, dict): 56 | return token_data.get("access_token") 57 | except (ValueError, KeyError): 58 | return None 59 | return None 60 | 61 | 62 | class RefreshingOauthToken: 63 | def __init__(self, hostname: str, scope: List[str], refresh_interval: int = 1800) -> None: 64 | self.hostname = hostname 65 | self.scope = scope 66 | self.refresh_interval = refresh_interval 67 | self.last_refresh_time = 0.0 68 | self.token = None 69 | 70 | def get_token(self) -> Any: 71 | current_time = time.time() 72 | if not self.token or current_time - self.last_refresh_time > self.refresh_interval: 73 | self.token = self._fetch_token() 74 | self.last_refresh_time = current_time 75 | return self.token 76 | 77 | def _fetch_token(self) -> Any: 78 | return oauth(self.hostname, self.scope) 79 | -------------------------------------------------------------------------------- /compute_modules/bin/ontology/generate_metadata_config.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import argparse 17 | import dataclasses 18 | import json 19 | from typing import Dict 20 | 21 | from compute_modules.bin.ontology._config_path import get_ontology_config_file 22 | from compute_modules.bin.ontology._types import ObjectTypeMetadata 23 | from compute_modules.bin.ontology.metadata_loader import load_object_type_metadata 24 | 25 | 26 | def write_inference_metadata( 27 | ontology_metadata: Dict[str, ObjectTypeMetadata], 28 | output_file: str, 29 | ) -> None: 30 | config = {"apiNameToTypeId": {key: value.objectTypeId for key, value in ontology_metadata.items()}} 31 | with open(output_file, "w") as f: 32 | json.dump(config, f) 33 | 34 | 35 | def main() -> None: 36 | parser = argparse.ArgumentParser() 37 | parser.add_argument( 38 | "--object-type-rid", 39 | required=False, 40 | help="Object type rid of ontology objects to import the metadata for", 41 | nargs="*", 42 | dest="object_type_rids", 43 | default=[], 44 | ) 45 | parser.add_argument( 46 | "--link-type-rid", 47 | required=False, 48 | help="Link type rid of ontology links to import the metadata for", 49 | nargs="*", 50 | dest="link_type_rids", 51 | default=[], 52 | ) 53 | parser.add_argument( 54 | "-t", 55 | "--token", 56 | required=True, 57 | help="Foundry token", 58 | default=None, 59 | ) 60 | parser.add_argument( 61 | "--foundry-url", 62 | required=True, 63 | help="Foundry stack url", 64 | default=None, 65 | ) 66 | parser.add_argument( 67 | "--ontology-metadata-config", 68 | required=False, 69 | help="Path to file for where to write the output configuration", 70 | dest="ontology_metadata_config_file", 71 | default=None, 72 | ) 73 | arguments = parser.parse_args() 74 | output_file = get_ontology_config_file(arguments.ontology_metadata_config_file) 75 | ontology_metadata = load_object_type_metadata( 76 | foundry_url=arguments.foundry_url, 77 | token=arguments.token, 78 | object_type_rids=arguments.object_type_rids, 79 | link_type_rids=arguments.link_type_rids, 80 | ) 81 | write_inference_metadata( 82 | ontology_metadata=ontology_metadata.objectMetadata, 83 | output_file=output_file, 84 | ) 85 | print(json.dumps(dataclasses.asdict(ontology_metadata))) 86 | 87 | 88 | if __name__ == "__main__": 89 | main() 90 | -------------------------------------------------------------------------------- /compute_modules/sources_v2/_api.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from dataclasses import dataclass, field 17 | from typing import Any, Optional 18 | 19 | # Env var constants 20 | SOURCE_CONFIGURATIONS_PATH = "SOURCE_CONFIGURATIONS_PATH" 21 | SOURCE_CREDENTIALS_PATH = "SOURCE_CREDENTIALS" 22 | SERVICE_DISCOVERY_PATH = "FOUNDRY_SERVICE_DISCOVERY_V2" 23 | DEFAULT_CA_BUNDLE = "DEFAULT_CA_PATH" 24 | 25 | 26 | JAVA_OFFSET_DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ" 27 | 28 | 29 | @dataclass 30 | class MountedClientCertificate: 31 | pem_certificate: str 32 | pem_private_key: str 33 | 34 | @staticmethod 35 | def from_dict(data: dict[str, Any]) -> "MountedClientCertificate": 36 | return MountedClientCertificate(pem_certificate=data["pemCertificate"], pem_private_key=data["pemPrivateKey"]) 37 | 38 | 39 | @dataclass 40 | class MountedHttpConnectionConfig: 41 | url: str 42 | auth_headers: dict[str, str] = field(default_factory=dict) 43 | query_parameters: dict[str, str] = field(default_factory=dict) 44 | 45 | @staticmethod 46 | def from_dict(data: dict[str, Any]) -> "MountedHttpConnectionConfig": 47 | return MountedHttpConnectionConfig( 48 | url=data["url"], auth_headers=data.get("authHeaders", {}), query_parameters=data.get("queryParameters", {}) 49 | ) 50 | 51 | 52 | @dataclass 53 | class MountedSourceConfig: 54 | secrets: dict[str, str] = field(default_factory=dict) 55 | http_connection_config: Optional[MountedHttpConnectionConfig] = None 56 | proxy_token: Optional[str] = None 57 | source_configuration: Any = None 58 | resolved_credentials: Any = None 59 | client_certificate: Optional[MountedClientCertificate] = None 60 | server_certificates: dict[str, str] = field(default_factory=dict) 61 | 62 | @staticmethod 63 | def from_dict(data: dict[str, Any]) -> "MountedSourceConfig": 64 | http_conn_config_data = data.get("httpConnectionConfig") 65 | http_conn_config = ( 66 | MountedHttpConnectionConfig.from_dict(http_conn_config_data) if http_conn_config_data is not None else None 67 | ) 68 | 69 | client_certificate_data = data.get("clientCertificate") 70 | client_certificate = ( 71 | MountedClientCertificate.from_dict(client_certificate_data) if client_certificate_data is not None else None 72 | ) 73 | 74 | return MountedSourceConfig( 75 | secrets=data.get("secrets", {}), 76 | proxy_token=data.get("proxyToken"), 77 | http_connection_config=http_conn_config, 78 | source_configuration=data.get("sourceConfiguration"), 79 | resolved_credentials=data.get("resolvedCredentials"), 80 | client_certificate=client_certificate, 81 | server_certificates=data.get("serverCertificates", {}), 82 | ) 83 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # Pycharm 156 | .idea 157 | 158 | # Visual Studio Code 159 | .vscode 160 | 161 | # Ruff 162 | .ruff_cache 163 | 164 | # Mac Explorer files 165 | .DS_Store -------------------------------------------------------------------------------- /scripts/ontology/metadata_loader.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from collections import defaultdict 17 | from typing import Dict, List 18 | 19 | from scripts.ontology._types import ObjectTypeMetadata, OntologyMetadataLinkTypeOuter, OntologyMetadataObjectTypeOuter 20 | from scripts.ontology.metadata_client import OntologyMetadataClient 21 | 22 | 23 | def load_object_type_metadata( 24 | foundry_url: str, 25 | token: str, 26 | object_type_rids: List[str], 27 | link_type_rids: List[str], 28 | ) -> Dict[str, ObjectTypeMetadata]: 29 | client = OntologyMetadataClient(foundry_url=foundry_url, token=token) 30 | object_types, link_types = client.bulk_load_entities( 31 | object_type_rids=object_type_rids, 32 | link_type_rids=link_type_rids, 33 | ) 34 | link_types_for_object = _get_link_types(link_types=link_types) 35 | return _get_object_type_metadata(object_types=object_types, link_types_for_object=link_types_for_object) 36 | 37 | 38 | def _get_object_type_metadata( 39 | object_types: List[OntologyMetadataObjectTypeOuter], 40 | link_types_for_object: Dict[str, Dict[str, str]], 41 | ) -> Dict[str, ObjectTypeMetadata]: 42 | object_type_metadata = {} 43 | for object_type in object_types: 44 | object_type_details = object_type["objectType"] 45 | primary_key_id = object_type_details["propertyTypes"][object_type_details["primaryKeys"][0]]["id"] 46 | object_rid = object_type_details["rid"] 47 | object_type_metadata[object_rid] = ObjectTypeMetadata( 48 | api_name=object_type_details["apiName"], 49 | type_id=object_type_details["id"], 50 | primary_key_id=primary_key_id, 51 | property_api_name_to_id={ 52 | value["apiName"]: value["id"] for _, value in object_type_details["propertyTypes"].items() 53 | }, 54 | link_type_api_name_to_id=link_types_for_object[object_rid], 55 | ) 56 | return object_type_metadata 57 | 58 | 59 | def _get_link_types(link_types: List[OntologyMetadataLinkTypeOuter]) -> Dict[str, Dict[str, str]]: 60 | link_types_for_object: Dict[str, Dict[str, str]] = defaultdict(dict) 61 | for link_type in link_types: 62 | link_type_id = link_type["linkType"]["id"] 63 | link_type_definition = link_type["linkType"]["definition"].get("manyToMany") 64 | # AFAIK oneToMany link edits are expressed as direct edits 65 | # on the ontology object so not implementing oneToMany for now 66 | if not link_type_definition: 67 | continue 68 | object_a_api_name = link_type_definition["objectTypeAToBLinkMetadata"].get("apiName") 69 | if object_a_api_name: 70 | link_types_for_object[link_type_definition["objectTypeRidA"]][object_a_api_name] = link_type_id 71 | object_b_api_name = link_type_definition["objectTypeBToALinkMetadata"].get("apiName") 72 | if object_b_api_name: 73 | link_types_for_object[link_type_definition["objectTypeRidB"]][object_b_api_name] = link_type_id 74 | return link_types_for_object 75 | -------------------------------------------------------------------------------- /compute_modules/bin/ontology/metadata_loader.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from collections import defaultdict 17 | from typing import Dict, List 18 | 19 | from compute_modules.bin.ontology._types import ( 20 | ObjectTypeMetadata, 21 | OntologyMetadataLinkTypeOuter, 22 | OntologyMetadataObjectTypeOuter, 23 | RuntimeMetadata, 24 | ) 25 | from compute_modules.bin.ontology.metadata_client import OntologyMetadataClient 26 | 27 | 28 | def load_object_type_metadata( 29 | foundry_url: str, 30 | token: str, 31 | object_type_rids: List[str], 32 | link_type_rids: List[str], 33 | ) -> RuntimeMetadata: 34 | client = OntologyMetadataClient(foundry_url=foundry_url, token=token) 35 | object_types, link_types = client.bulk_load_entities( 36 | object_type_rids=object_type_rids, 37 | link_type_rids=link_type_rids, 38 | ) 39 | link_types_for_object = _get_link_types(link_types=link_types) 40 | object_type_metadata = _get_object_type_metadata( 41 | object_types=object_types, link_types_for_object=link_types_for_object 42 | ) 43 | return RuntimeMetadata(objectMetadata=object_type_metadata) 44 | 45 | 46 | def _get_object_type_metadata( 47 | object_types: List[OntologyMetadataObjectTypeOuter], 48 | link_types_for_object: Dict[str, Dict[str, str]], 49 | ) -> Dict[str, ObjectTypeMetadata]: 50 | object_type_metadata = {} 51 | for object_type in object_types: 52 | object_type_details = object_type["objectType"] 53 | if object_type_details["apiName"] in object_type_metadata: 54 | raise ValueError(f"Duplicate api name found in ontology metadata: {object_type_details['apiName']}") 55 | primary_key_id = object_type_details["propertyTypes"][object_type_details["primaryKeys"][0]]["id"] 56 | object_rid = object_type_details["rid"] 57 | object_type_metadata[object_type_details["apiName"]] = ObjectTypeMetadata( 58 | objectTypeApiName=object_type_details["apiName"], 59 | objectTypeId=object_type_details["id"], 60 | primaryKeyPropertyId=primary_key_id, 61 | properties={value["apiName"]: value["id"] for _, value in object_type_details["propertyTypes"].items()}, 62 | links=link_types_for_object[object_rid], 63 | ) 64 | return object_type_metadata 65 | 66 | 67 | def _get_link_types(link_types: List[OntologyMetadataLinkTypeOuter]) -> Dict[str, Dict[str, str]]: 68 | link_types_for_object: Dict[str, Dict[str, str]] = defaultdict(dict) 69 | for link_type in link_types: 70 | link_type_id = link_type["linkType"]["id"] 71 | link_type_definition = link_type["linkType"]["definition"].get("manyToMany") 72 | # AFAIK oneToMany link edits are expressed as direct edits 73 | # on the ontology object so not implementing oneToMany for now 74 | if not link_type_definition: 75 | continue 76 | object_a_api_name = link_type_definition["objectTypeAToBLinkMetadata"].get("apiName") 77 | if object_a_api_name: 78 | link_types_for_object[link_type_definition["objectTypeRidA"]][object_a_api_name] = link_type_id 79 | object_b_api_name = link_type_definition["objectTypeBToALinkMetadata"].get("apiName") 80 | if object_b_api_name: 81 | link_types_for_object[link_type_definition["objectTypeRidB"]][object_b_api_name] = link_type_id 82 | return link_types_for_object 83 | -------------------------------------------------------------------------------- /compute_modules/sources/_sources.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import json 17 | import os 18 | import warnings 19 | from dataclasses import dataclass, field 20 | from typing import Any, Dict, Optional 21 | 22 | # Global Mutable State 23 | _source_credentials = None 24 | _source_configurations = None 25 | 26 | # Env var constants 27 | SOURCE_CONFIGURATIONS_PATH = "SOURCE_CONFIGURATIONS_PATH" 28 | SOURCE_CREDENTIALS_PATH = "SOURCE_CREDENTIALS" 29 | 30 | 31 | @dataclass 32 | class MountedHttpConnectionConfig: 33 | url: str 34 | auth_headers: Dict[str, str] = field(default_factory=dict) 35 | query_parameters: Dict[str, str] = field(default_factory=dict) 36 | 37 | @staticmethod 38 | def from_dict(data: Dict[str, Any]) -> "MountedHttpConnectionConfig": 39 | return MountedHttpConnectionConfig( 40 | url=data["url"], auth_headers=data.get("authHeaders", {}), query_parameters=data.get("queryParameters", {}) 41 | ) 42 | 43 | 44 | @dataclass 45 | class MountedSourceConfig: 46 | secrets: Dict[str, str] = field(default_factory=dict) 47 | http_connection_config: Optional[MountedHttpConnectionConfig] = None 48 | proxy_token: Optional[str] = None 49 | source_configuration: Any = None 50 | resolved_credentials: Any = None 51 | 52 | @staticmethod 53 | def from_dict(data: Dict[str, Any]) -> "MountedSourceConfig": 54 | http_conn_config_data = data.get("httpConnectionConfig") 55 | http_conn_config = ( 56 | MountedHttpConnectionConfig.from_dict(http_conn_config_data) if http_conn_config_data is not None else None 57 | ) 58 | return MountedSourceConfig( 59 | secrets=data.get("secrets", {}), 60 | proxy_token=data.get("proxyToken"), 61 | http_connection_config=http_conn_config, 62 | source_configuration=data.get("sourceConfiguration"), 63 | resolved_credentials=data.get("resolvedCredentials"), 64 | ) 65 | 66 | 67 | def get_sources() -> Dict[str, Dict[str, str]]: 68 | warnings.warn( 69 | "get_sources is deprecated. Use get_source in compute_modules.sources_v2 instead.", 70 | DeprecationWarning, 71 | stacklevel=2, 72 | ) 73 | global _source_credentials 74 | if _source_credentials is None: 75 | creds_path = os.environ.get(SOURCE_CREDENTIALS_PATH) 76 | if creds_path: 77 | with open(creds_path, "r", encoding="utf-8") as fr: 78 | data = json.load(fr) 79 | if isinstance(data, dict): 80 | _source_credentials = data 81 | else: 82 | raise ValueError("The JSON content is not a dictionary") 83 | return _source_credentials if _source_credentials is not None else {} 84 | 85 | 86 | def get_source_configurations() -> Dict[str, Any]: 87 | warnings.warn( 88 | "get_source_configurations is deprecated. Use get_source in compute_modules.sources_v2 instead.", 89 | DeprecationWarning, 90 | stacklevel=2, 91 | ) 92 | global _source_configurations 93 | if _source_configurations is None: 94 | configs_path = os.environ.get(SOURCE_CONFIGURATIONS_PATH) 95 | if configs_path: 96 | with open(configs_path, "r", encoding="utf-8") as fr: 97 | raw_configs = json.load(fr) 98 | if isinstance(raw_configs, dict): 99 | _source_configurations = { 100 | key: MountedSourceConfig.from_dict(value) for key, value in raw_configs.items() 101 | } 102 | else: 103 | raise ValueError("The JSON content is not a dictionary") 104 | configs = _source_configurations if _source_configurations is not None else {} 105 | return {key: value.source_configuration for key, value in configs.items()} 106 | 107 | 108 | def get_source_secret(source_api_name: str, credential_name: str) -> Any: 109 | warnings.warn( 110 | "get_source_secret is deprecated. Use get_source in compute_modules.sources_v2 instead.", 111 | DeprecationWarning, 112 | stacklevel=2, 113 | ) 114 | source_credentials = get_sources() 115 | return source_credentials.get(source_api_name, {}).get(credential_name) 116 | 117 | 118 | def get_source_config(source_api_name: str) -> Any: 119 | warnings.warn( 120 | "get_source_config is deprecated. Use get_source in compute_modules.sources_v2 instead.", 121 | DeprecationWarning, 122 | stacklevel=2, 123 | ) 124 | return get_source_configurations().get(source_api_name, {}) 125 | -------------------------------------------------------------------------------- /compute_modules/startup.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import os 17 | import threading 18 | from enum import Enum 19 | from multiprocessing.pool import Pool, ThreadPool 20 | from typing import Any, Dict, Optional, Union 21 | 22 | from compute_modules.client.internal_query_client import InternalQueryService 23 | from compute_modules.function_registry.function_registry import ( 24 | FUNCTION_SCHEMA_CONVERSIONS, 25 | FUNCTION_SCHEMAS, 26 | IS_FUNCTION_CONTEXT_TYPED, 27 | REGISTERED_FUNCTIONS, 28 | STREAMING, 29 | ) 30 | 31 | # This is a workaround to prevent starting the CM when doing static function inference at build time 32 | DISABLE_STARTUP = False 33 | QUERY_CLIENT: Optional[InternalQueryService] = None 34 | 35 | 36 | class ConcurrencyType(str, Enum): 37 | PROCESS_POOL = "PROCESS_POOL" 38 | THREAD_POOL = "THREAD_POOL" 39 | 40 | 41 | # The main process creates the initial InternalQueryService instance, 42 | # which is used to post the function schema and poll for jobs. 43 | # 44 | # We then spin up a Pool of worker processes that all create their own sessions. 45 | # When a job is received by the main process, that job is delegated to a worker. 46 | # The worker then processes that job and posts the result back to the runtime using its own session. 47 | # 48 | # This library supports streaming responses via generators. 49 | # Python passes data between processes by using `pickle` to serialize the data. 50 | # Generators cannot be pickled, so we cannot pass the result of streaming functions back to the parent process. 51 | # As a workaround, the worker process is given a session only for posting job results back to the runtime. 52 | 53 | 54 | def _handle_job(job: Dict[str, Any]) -> None: 55 | """Helper function to be called by pool.apply_async since python can't serialize methods""" 56 | global QUERY_CLIENT 57 | assert QUERY_CLIENT, "QUERY_CLIENT is uninitialized" 58 | QUERY_CLIENT.handle_job(job=job) 59 | 60 | 61 | def _get_and_schedule_job(pool: Union[Pool, ThreadPool]) -> None: 62 | """Try to get a job and schedule to the process pool""" 63 | global QUERY_CLIENT 64 | assert QUERY_CLIENT, "QUERY_CLIENT is uninitialized" 65 | job = None 66 | try: 67 | job = QUERY_CLIENT.get_job_or_none() 68 | except Exception as e: 69 | QUERY_CLIENT.logger.warning(f"Exception occurred while fetching job: {str(e)}") 70 | if job: 71 | QUERY_CLIENT.logger.debug(f"Got a job: {job}") 72 | pool.apply_async(_handle_job, (job,)) 73 | 74 | 75 | def _worker_process_init() -> None: 76 | """Create a new session for each worker process""" 77 | global QUERY_CLIENT 78 | assert QUERY_CLIENT, "QUERY_CLIENT is uninitialized" 79 | QUERY_CLIENT.init_session() 80 | QUERY_CLIENT._set_logger_process_id(os.getpid()) 81 | 82 | 83 | def _worker_thread_init() -> None: 84 | """Create a new session for each worker thread""" 85 | global QUERY_CLIENT 86 | assert QUERY_CLIENT, "QUERY_CLIENT is uninitialized" 87 | QUERY_CLIENT._set_logger_process_id(threading.get_ident()) 88 | 89 | 90 | def start_compute_module( 91 | concurrency_type: ConcurrencyType = ConcurrencyType.PROCESS_POOL, 92 | report_restart: bool = True, 93 | ) -> None: 94 | """Starts a Compute Module that will Poll for jobs indefinitely""" 95 | if DISABLE_STARTUP: 96 | return 97 | 98 | global QUERY_CLIENT 99 | QUERY_CLIENT = InternalQueryService( 100 | registered_functions=REGISTERED_FUNCTIONS, 101 | function_schemas=FUNCTION_SCHEMAS, 102 | function_schema_conversions=FUNCTION_SCHEMA_CONVERSIONS, 103 | is_function_context_typed=IS_FUNCTION_CONTEXT_TYPED, 104 | streaming=STREAMING, 105 | ) 106 | QUERY_CLIENT.post_query_schemas() 107 | if report_restart: 108 | QUERY_CLIENT.report_restart() 109 | QUERY_CLIENT.logger.info(f"Starting to poll for jobs with concurrency {QUERY_CLIENT.concurrency}") 110 | if concurrency_type == ConcurrencyType.PROCESS_POOL: 111 | with Pool(QUERY_CLIENT.concurrency, initializer=_worker_process_init) as pool: 112 | while True: 113 | QUERY_CLIENT.logger.info("Polling for new jobs...") 114 | _get_and_schedule_job(pool) 115 | else: 116 | with ThreadPool(QUERY_CLIENT.concurrency, initializer=_worker_thread_init) as pool: 117 | while True: 118 | QUERY_CLIENT.logger.info("Polling for new jobs...") 119 | _get_and_schedule_job(pool) 120 | 121 | 122 | __all__ = [ 123 | "ConcurrencyType", 124 | "start_compute_module", 125 | ] 126 | -------------------------------------------------------------------------------- /tests/logging/test_log_levels.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import logging 17 | 18 | import pytest 19 | 20 | from compute_modules.logging import get_logger, set_internal_log_level 21 | 22 | from .logging_test_utils import ( 23 | CLIENT_CRITICAL_STR, 24 | CLIENT_DEBUG_STR, 25 | CLIENT_ERROR_STR, 26 | CLIENT_INFO_STR, 27 | CLIENT_WARNING_STR, 28 | CRITICAL_STR, 29 | DEBUG_STR, 30 | ERROR_STR, 31 | INFO_STR, 32 | WARNING_STR, 33 | emit_internal_logs, 34 | ) 35 | 36 | 37 | def test_default_log_levels( 38 | caplog: pytest.LogCaptureFixture, 39 | ) -> None: 40 | """Test default log level is ERROR""" 41 | # Need to actually re-set it here since we're overriding in other tests 42 | set_internal_log_level(logging.ERROR) 43 | emit_internal_logs() 44 | logged_messages = [record.message for record in caplog.records] 45 | 46 | assert DEBUG_STR not in logged_messages 47 | assert INFO_STR not in logged_messages 48 | assert WARNING_STR not in logged_messages 49 | assert ERROR_STR in logged_messages 50 | assert CRITICAL_STR in logged_messages 51 | 52 | 53 | def test_log_level_override( 54 | caplog: pytest.LogCaptureFixture, 55 | ) -> None: 56 | """Test calling set_internal_log_level to change CM log level""" 57 | set_internal_log_level(logging.INFO) 58 | emit_internal_logs() 59 | assert DEBUG_STR not in caplog.text 60 | assert INFO_STR in caplog.text 61 | assert WARNING_STR in caplog.text 62 | assert ERROR_STR in caplog.text 63 | assert CRITICAL_STR in caplog.text 64 | 65 | 66 | def test_log_level_override_with_client_level_lower( 67 | caplog: pytest.LogCaptureFixture, 68 | ) -> None: 69 | """Test calling set_internal_log_level to change CM log level, 70 | while having public logger at a different level 71 | """ 72 | set_internal_log_level(logging.WARNING) 73 | client_logger = get_logger("twinkle") 74 | client_logger.setLevel(logging.DEBUG) 75 | client_logger.debug(CLIENT_DEBUG_STR) 76 | client_logger.info(CLIENT_INFO_STR) 77 | client_logger.warning(CLIENT_WARNING_STR) 78 | client_logger.error(CLIENT_ERROR_STR) 79 | client_logger.critical(CLIENT_CRITICAL_STR) 80 | emit_internal_logs() 81 | assert DEBUG_STR not in caplog.text 82 | assert INFO_STR not in caplog.text 83 | assert WARNING_STR in caplog.text 84 | assert ERROR_STR in caplog.text 85 | assert CRITICAL_STR in caplog.text 86 | 87 | assert CLIENT_DEBUG_STR in caplog.text 88 | assert CLIENT_INFO_STR in caplog.text 89 | assert CLIENT_WARNING_STR in caplog.text 90 | assert CLIENT_ERROR_STR in caplog.text 91 | assert CLIENT_CRITICAL_STR in caplog.text 92 | 93 | 94 | def test_log_level_override_with_client_level_higher( 95 | caplog: pytest.LogCaptureFixture, 96 | ) -> None: 97 | """Test calling set_internal_log_level to change CM log level, 98 | while having public logger at a different level 99 | """ 100 | set_internal_log_level(logging.DEBUG) 101 | client_logger = get_logger("twinkle") 102 | client_logger.setLevel(logging.ERROR) 103 | client_logger.debug(CLIENT_DEBUG_STR) 104 | client_logger.info(CLIENT_INFO_STR) 105 | client_logger.warning(CLIENT_WARNING_STR) 106 | client_logger.error(CLIENT_ERROR_STR) 107 | client_logger.critical(CLIENT_CRITICAL_STR) 108 | emit_internal_logs() 109 | assert DEBUG_STR in caplog.text 110 | assert INFO_STR in caplog.text 111 | assert WARNING_STR in caplog.text 112 | assert ERROR_STR in caplog.text 113 | assert CRITICAL_STR in caplog.text 114 | 115 | assert CLIENT_DEBUG_STR not in caplog.text 116 | assert CLIENT_INFO_STR not in caplog.text 117 | assert CLIENT_WARNING_STR not in caplog.text 118 | assert CLIENT_ERROR_STR in caplog.text 119 | assert CLIENT_CRITICAL_STR in caplog.text 120 | assert CRITICAL_STR in caplog.text 121 | 122 | assert CLIENT_DEBUG_STR not in caplog.text 123 | assert CLIENT_INFO_STR not in caplog.text 124 | assert CLIENT_WARNING_STR not in caplog.text 125 | assert CLIENT_ERROR_STR in caplog.text 126 | assert CLIENT_CRITICAL_STR in caplog.text 127 | assert CLIENT_CRITICAL_STR in caplog.text 128 | assert CLIENT_INFO_STR not in caplog.text 129 | assert CLIENT_WARNING_STR not in caplog.text 130 | assert CLIENT_ERROR_STR in caplog.text 131 | assert CLIENT_CRITICAL_STR in caplog.text 132 | assert CRITICAL_STR in caplog.text 133 | 134 | assert CLIENT_DEBUG_STR not in caplog.text 135 | assert CLIENT_INFO_STR not in caplog.text 136 | assert CLIENT_WARNING_STR not in caplog.text 137 | assert CLIENT_ERROR_STR in caplog.text 138 | assert CLIENT_CRITICAL_STR in caplog.text 139 | assert CLIENT_CRITICAL_STR in caplog.text 140 | -------------------------------------------------------------------------------- /tests/function_registry/test_function_payload_converter.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import datetime 17 | import decimal 18 | 19 | import pytest 20 | 21 | from compute_modules.function_registry.function_payload_converter import convert_payload 22 | from compute_modules.function_registry.function_schema_parser import parse_function_schema 23 | from tests.function_registry.dummy_app import ChildClass, DummyInput, ParentClass, dummy_func_1 24 | 25 | RAW_PAYLOAD = { 26 | "parent_class": { 27 | "some_flag": False, 28 | "some_value": 1234, 29 | "child": {"timestamp": "1979-05-27T07:32:00Z", "some_value": 1.234, "another_optional_field": "something!"}, 30 | }, 31 | "optional_field": None, 32 | "set_field": ["2024-09-04", "2024-07-20", "1984-05-19"], 33 | "map_field": {"dmFsdWU=": "1.0", "dmFsdWUy": "2.0"}, 34 | "some_flag": True, 35 | "datetime_list": ["2024-09-04T12:00:00Z", "2024-07-20T15:30:00Z", "1984-05-19T08:45:00Z"], 36 | } 37 | 38 | BAD_RAW_PAYLOAD = { 39 | "parent_class": { 40 | "some_flag": False, 41 | "some_value": 1234, 42 | "child": {"timestamp": "1979-05-27T07:32:00Z", "some_value": 1.234, "another_optional_field": "something!"}, 43 | }, 44 | "optional_field": None, 45 | "set_field": ["do", "re", "mi"], 46 | "map_field": {"dmFsdWU=": "1.0", "dmFsdWUy": "2.0"}, 47 | "some_flag": True, 48 | "datetime_list": ["2024-09-04T12:00:00Z", "2024-07-20T15:30:00Z", "1984-05-19T08:45:00Z"], 49 | } 50 | 51 | 52 | @pytest.fixture() 53 | def expected_return_value() -> DummyInput: 54 | parent_dict = RAW_PAYLOAD["parent_class"] 55 | child_dict = parent_dict["child"] # type: ignore[index, call-overload] 56 | child = ChildClass( 57 | timestamp=datetime.datetime.strptime(child_dict["timestamp"], "%Y-%m-%dT%H:%M:%SZ"), # type: ignore[index] 58 | some_value=child_dict["some_value"], # type: ignore[index] 59 | another_optional_field=child_dict["another_optional_field"], # type: ignore[index] 60 | ) 61 | parent = ParentClass( 62 | some_flag=parent_dict["some_flag"], # type: ignore[index, call-overload, arg-type] 63 | some_value=parent_dict["some_value"], # type: ignore[index, call-overload, arg-type] 64 | child=child, 65 | ) 66 | datetime_list = [datetime.datetime.strptime(dt, "%Y-%m-%dT%H:%M:%SZ") for dt in RAW_PAYLOAD["datetime_list"]] # type: ignore[union-attr] 67 | set_field = set([datetime.date.fromisoformat(d) for d in RAW_PAYLOAD["set_field"]]) # type: ignore[union-attr] 68 | map_field = {bytes(k, encoding="utf8"): decimal.Decimal(v) for k, v in RAW_PAYLOAD["map_field"].items()} # type: ignore[union-attr, arg-type] 69 | return DummyInput( 70 | parent_class=parent, 71 | optional_field=RAW_PAYLOAD["optional_field"], # type: ignore[arg-type] 72 | set_field=set_field, 73 | datetime_list=datetime_list, 74 | map_field=map_field, 75 | some_flag=RAW_PAYLOAD["some_flag"], # type: ignore[arg-type] 76 | ) 77 | 78 | 79 | def test_convert_payload( 80 | expected_return_value: DummyInput, 81 | ) -> None: 82 | """Test the happy path for convert_payload""" 83 | parse_result = parse_function_schema(dummy_func_1, "dummy_func_1", [], {}) 84 | assert parse_result.class_node 85 | processed_payload: DummyInput = convert_payload(RAW_PAYLOAD, parse_result.class_node) 86 | assert processed_payload is not None 87 | assert processed_payload.parent_class.child.__dict__ == expected_return_value.parent_class.child.__dict__ 88 | assert processed_payload.parent_class.some_flag == expected_return_value.parent_class.some_flag 89 | assert processed_payload.parent_class.some_value == expected_return_value.parent_class.some_value 90 | assert processed_payload.optional_field == expected_return_value.optional_field 91 | assert processed_payload.set_field == expected_return_value.set_field 92 | assert processed_payload.map_field == expected_return_value.map_field 93 | assert processed_payload.some_flag == expected_return_value.some_flag 94 | assert processed_payload.optional_default_value_field == expected_return_value.optional_default_value_field 95 | assert processed_payload.datetime_list == expected_return_value.datetime_list 96 | 97 | 98 | def test_convert_payload_error( 99 | caplog: pytest.LogCaptureFixture, 100 | ) -> None: 101 | """Test the error path for convert_payload""" 102 | parse_result = parse_function_schema(dummy_func_1, "dummy_func_1", [], {}) 103 | assert parse_result.class_node 104 | with pytest.raises(ValueError) as exc_info: 105 | convert_payload(BAD_RAW_PAYLOAD, parse_result.class_node) 106 | assert str(exc_info.value) == "Invalid isoformat string: 'do'" 107 | assert "Error converting do to type tuple[ComputeModulesLoggerAdapter, ComputeModulesLoggerAdapter, ComputeModulesLoggerAdapter]: 32 | """Initializes & configures loggers""" 33 | 34 | internal_logger = internal.get_internal_logger() 35 | internal_logger.setLevel(logging.INFO) 36 | logger_1 = get_logger("test.logger.1") 37 | logger_1.setLevel(logging.INFO) 38 | logger_2 = get_logger("test.logger.2") 39 | logger_2.setLevel(logging.INFO) 40 | return (internal_logger, logger_1, logger_2) 41 | 42 | 43 | def format_log_context(pid: int, job_id: str) -> str: 44 | return f"PID: {pid:<6} JOB: {job_id:<37}" 45 | 46 | 47 | def test_log_format(capsys: pytest.CaptureFixture[str], custom_formatter: JsonFormatter) -> None: 48 | """Verify initial state of logger context""" 49 | internal_logger, logger_1, logger_2 = logger_fixtures() 50 | internal_logger.info(INFO_STR) 51 | logger_1.info(CLIENT_INFO_STR) 52 | logger_2.info(CLIENT_WARNING_STR) 53 | captured = capsys.readouterr() 54 | parsed_out = list(filter(lambda x: x, captured.err.split("\n"))) 55 | assert len(parsed_out) == 3 56 | for log in parsed_out: 57 | assert format_log_context(pid=-1, job_id="") in log 58 | 59 | COMPUTE_MODULES_ADAPTER_MANAGER.update_process_id(process_id=PROCESS_ID) 60 | internal_logger.info(INFO_STR) 61 | logger_1.info(CLIENT_INFO_STR) 62 | logger_2.info(CLIENT_WARNING_STR) 63 | captured = capsys.readouterr() 64 | parsed_out = list(filter(lambda x: x, captured.err.split("\n"))) 65 | assert len(parsed_out) == 3 66 | for log in parsed_out: 67 | assert format_log_context(pid=PROCESS_ID, job_id="") in log 68 | 69 | job_id = str(uuid.uuid4()) 70 | COMPUTE_MODULES_ADAPTER_MANAGER.update_job_id(job_id=job_id) 71 | internal_logger.info(INFO_STR) 72 | logger_1.info(CLIENT_INFO_STR) 73 | logger_2.info(CLIENT_WARNING_STR) 74 | captured = capsys.readouterr() 75 | parsed_out = list(filter(lambda x: x, captured.err.split("\n"))) 76 | assert len(parsed_out) == 3 77 | for log in parsed_out: 78 | assert format_log_context(pid=PROCESS_ID, job_id=job_id) in log 79 | 80 | # Test clearing now 81 | COMPUTE_MODULES_ADAPTER_MANAGER.update_job_id(job_id="") 82 | internal_logger.info(INFO_STR) 83 | logger_1.info(CLIENT_INFO_STR) 84 | logger_2.info(CLIENT_WARNING_STR) 85 | captured = capsys.readouterr() 86 | parsed_out = list(filter(lambda x: x, captured.err.split("\n"))) 87 | assert len(parsed_out) == 3 88 | for log in parsed_out: 89 | assert format_log_context(pid=PROCESS_ID, job_id="") in log 90 | 91 | # Test Custom Formatting 92 | # TODO split out into custom test once ComputeModuleLoggingAdapter made fixture to avoid capsys errors 93 | 94 | client_logger = get_logger("twinkle") 95 | client_logger.setLevel(logging.INFO) 96 | setup_logger_formatter(custom_formatter) 97 | 98 | job_id = str(uuid.uuid4()) 99 | process_id = 5 100 | 101 | COMPUTE_MODULES_ADAPTER_MANAGER.update_process_id(process_id) 102 | COMPUTE_MODULES_ADAPTER_MANAGER.update_job_id(job_id) 103 | 104 | client_logger.info(CLIENT_INFO_STR) 105 | 106 | logged = capsys.readouterr().err 107 | 108 | try: 109 | log_js = json.loads(logged) 110 | valid_json = True 111 | except ValueError: 112 | valid_json = False 113 | 114 | # Test external logger 115 | assert valid_json, "Logs should be json" 116 | assert log_js["level"] == "INFO", "Log has wrong level" 117 | assert log_js["message"] == CLIENT_INFO_STR, "Log has wrong message" 118 | assert log_js["job_id"] == job_id 119 | assert log_js["process_id"] == str(process_id) 120 | 121 | internal_logger = internal.get_internal_logger() 122 | 123 | internal_logger.error(CLIENT_ERROR_STR) 124 | internal_logged = capsys.readouterr().err 125 | 126 | try: 127 | interal_log_js = json.loads(internal_logged) 128 | internal_valid_json = True 129 | except ValueError: 130 | internal_valid_json = False 131 | 132 | # Test internal logger 133 | assert internal_valid_json, "Logs should be json" 134 | assert interal_log_js["level"] == "ERROR", "Log has wrong level" 135 | assert interal_log_js["message"] == CLIENT_ERROR_STR, "Log has wrong message" 136 | assert interal_log_js["job_id"] == job_id 137 | assert interal_log_js["process_id"] == str(process_id) 138 | -------------------------------------------------------------------------------- /compute_modules/logging/common.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import json 17 | import logging 18 | import threading 19 | from datetime import datetime, timezone 20 | from typing import TYPE_CHECKING, Any, Dict, MutableMapping, Optional, Tuple, Union 21 | 22 | # logging.LoggerAdapter was made generic in 3.11 so we need to determine at runtime 23 | # whether this should be generic or not. 24 | # 25 | # See: https://mypy.readthedocs.io/en/stable/runtime_troubles.html#using-classes-that-are-generic-in-stubs-but-not-at-runtime 26 | # 27 | if TYPE_CHECKING: 28 | _LoggerAdapter = logging.LoggerAdapter[logging.Logger] 29 | else: 30 | _LoggerAdapter = logging.LoggerAdapter 31 | 32 | DEFAULT_LOG_FORMAT = "PID: %(process_id)-6s JOB: %(job_id)-36s - %(message)s" 33 | DEFAULT_LOG_STRING_FORMATTER = logging.Formatter(DEFAULT_LOG_FORMAT) 34 | 35 | 36 | class SlsFormatter(logging.Formatter): 37 | """Custom SLS formatter for structured logging by sidecar""" 38 | 39 | def __init__(self) -> None: 40 | super().__init__() 41 | self.string_formatter = DEFAULT_LOG_STRING_FORMATTER 42 | 43 | def format(self, record: Any) -> str: 44 | # Use the default string formatter for the message field 45 | formatted_message = self.string_formatter.format(record) 46 | 47 | log_entry = { 48 | "type": getattr(record, "service_type", "service.1"), 49 | "level": record.levelname, 50 | "time": datetime.now(timezone.utc).isoformat(), 51 | "origin": f"{record.filename}:{record.lineno}", 52 | "safe": True, 53 | "thread": threading.current_thread().name, 54 | "message": formatted_message, 55 | } 56 | return json.dumps(log_entry) 57 | 58 | 59 | SLS_FORMATTER = SlsFormatter() 60 | LOG_FORMATTER = None 61 | 62 | 63 | def _setup_logger_formatter( 64 | formatter: logging.Formatter, 65 | ) -> None: 66 | if formatter: 67 | global LOG_FORMATTER 68 | LOG_FORMATTER = formatter 69 | 70 | for adapter in COMPUTE_MODULES_ADAPTER_MANAGER.adapters.values(): 71 | for handler in adapter.logger.handlers: 72 | handler.setFormatter(LOG_FORMATTER) 73 | 74 | 75 | # TODO: support for log file output (need access to selected log output location) 76 | def _create_logger(name: str) -> logging.Logger: 77 | """Creates a logger that can have its log level set ... and actually work. 78 | 79 | See: https://stackoverflow.com/a/59705351 80 | """ 81 | logger = logging.getLogger(name) 82 | handler = logging.StreamHandler() 83 | formatter = LOG_FORMATTER if LOG_FORMATTER else SLS_FORMATTER 84 | handler.setFormatter(formatter) 85 | logger.handlers.clear() 86 | logger.addHandler(handler) 87 | 88 | return logger 89 | 90 | 91 | THREAD_LOCAL = threading.local() 92 | 93 | 94 | def set_thread_local_data(key: str, value: str) -> None: 95 | setattr(THREAD_LOCAL, key, value) 96 | 97 | 98 | def get_thread_local_data(key: str, default: str) -> str: 99 | return getattr(THREAD_LOCAL, key, default) 100 | 101 | 102 | # Custom LoggerAdapter to inject job- & thread/process-specific information into log lines 103 | # 104 | # See: https://docs.python.org/3/howto/logging-cookbook.html#using-loggeradapters-to-impart-contextual-information 105 | class ComputeModulesLoggerAdapter(_LoggerAdapter): 106 | """Wrapper around Python's `logging.LoggerAdapter` class. 107 | This can be used like a normal `logging.Logger` instance 108 | """ 109 | 110 | def __init__( 111 | self, 112 | logger_name: str, 113 | ) -> None: 114 | # Need to pass empty dict as `extra` param for 3.9 support 115 | super().__init__(_create_logger(logger_name), dict()) 116 | 117 | def process(self, msg: Any, kwargs: MutableMapping[str, Any]) -> Tuple[Any, MutableMapping[str, Any]]: 118 | custom_data = { 119 | "process_id": str(get_thread_local_data("process_id", "-1")), 120 | "job_id": str(get_thread_local_data("job_id", "")), 121 | } 122 | kwargs["extra"] = kwargs.get("extra", {}) 123 | kwargs["extra"].update(custom_data) 124 | 125 | return msg, kwargs 126 | 127 | 128 | class ComputeModulesAdapterManager(object): 129 | adapters: Dict[str, ComputeModulesLoggerAdapter] = {} 130 | 131 | def get_logger(self, name: str, default_level: Optional[Union[str, int]] = None) -> ComputeModulesLoggerAdapter: 132 | """Get a logger by name. If it does not already exist, creates it first""" 133 | if name not in self.adapters: 134 | self.adapters[name] = ComputeModulesLoggerAdapter(name) 135 | if default_level: 136 | self.adapters[name].setLevel(default_level) 137 | return self.adapters[name] 138 | 139 | def update_process_id(self, process_id: int) -> None: 140 | """Update process_id for all registered adapters""" 141 | set_thread_local_data("process_id", str(process_id)) 142 | 143 | def update_job_id(self, job_id: str) -> None: 144 | """Update job_id for all registered adapters""" 145 | set_thread_local_data("job_id", str(job_id)) 146 | 147 | 148 | COMPUTE_MODULES_ADAPTER_MANAGER = ComputeModulesAdapterManager() 149 | 150 | 151 | __all__ = [ 152 | "COMPUTE_MODULES_ADAPTER_MANAGER", 153 | "ComputeModulesLoggerAdapter", 154 | "_setup_logger_formatter", 155 | ] 156 | -------------------------------------------------------------------------------- /scripts/checks.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import shutil 17 | import subprocess 18 | import sys 19 | from datetime import date 20 | from itertools import chain 21 | from os import unlink 22 | from pathlib import Path 23 | from typing import Iterator, List, Tuple 24 | 25 | SOURCE_DIR = "compute_modules" 26 | TESTS_DIR = "tests" 27 | SCRIPTS_DIR = "scripts" 28 | LICENSE_FILE = "LICENSE.txt" 29 | FILES_WITH_LICENSE_NEEDED_GLOB_EXPR = "*.py" 30 | 31 | 32 | def test() -> None: 33 | """Runs pytest on all tests in the tests/ directory""" 34 | result = subprocess.run(["pytest", TESTS_DIR]) 35 | sys.exit(result.returncode) 36 | 37 | 38 | def check_format() -> None: 39 | """Runs linter 'checks'. Raises exception if any linting exceptions exist""" 40 | black_result = subprocess.run(["black", "--check", "--diff", SOURCE_DIR, TESTS_DIR, SCRIPTS_DIR]) 41 | ruff_result = subprocess.run(["ruff", "check", SOURCE_DIR, TESTS_DIR, SCRIPTS_DIR]) 42 | isort_result = subprocess.run(["isort", "--check", "--diff", SOURCE_DIR, TESTS_DIR, SCRIPTS_DIR]) 43 | exit_code = 1 if black_result.returncode or ruff_result.returncode or isort_result.returncode else 0 44 | sys.exit(exit_code) 45 | 46 | 47 | def check_mypy() -> None: 48 | """Runs mypy checks. Raises exception if any mypy exceptions exist""" 49 | result = subprocess.run(["mypy", SOURCE_DIR, TESTS_DIR, SCRIPTS_DIR]) 50 | sys.exit(result.returncode) 51 | 52 | 53 | def format() -> None: 54 | """Formats all files to fix any linter issues that can be automatically fixed""" 55 | subprocess.run(["black", SOURCE_DIR, TESTS_DIR, SCRIPTS_DIR]) 56 | subprocess.run(["ruff", "format", SOURCE_DIR, TESTS_DIR, SCRIPTS_DIR]) 57 | subprocess.run(["isort", SOURCE_DIR, TESTS_DIR, SCRIPTS_DIR]) 58 | 59 | 60 | def _get_license_content( 61 | year: int, 62 | ) -> Tuple[str, int]: 63 | with open(LICENSE_FILE, "r", encoding="utf-8") as f: 64 | lines = [] 65 | for line in f.readlines(): 66 | # Empty lines should not have whitespace appended 67 | if line == "\n": 68 | lines.append("#\n") 69 | else: 70 | lines.append(f"# {line}") 71 | content = "".join(lines) 72 | content = content.format(YEAR=year) 73 | return content, len(lines) 74 | 75 | 76 | def _get_license_content_any_year() -> Tuple[List[str], int]: 77 | start_year = 2024 78 | current_year = date.today().year 79 | num_lines = -1 80 | contents = [] 81 | for year in range(start_year, current_year + 1): 82 | content, num_lines = _get_license_content(year) 83 | contents.append(content) 84 | return contents, num_lines 85 | 86 | 87 | def _get_n_lines_of_file(filename: str, num_lines: int) -> str: 88 | try: 89 | with open(filename, "r", encoding="utf-8") as f: 90 | head = [] 91 | for _ in range(num_lines): 92 | head.append(next(f)) 93 | except StopIteration: 94 | # File doesn't have as many lines as license file 95 | pass 96 | return "".join(head) 97 | 98 | 99 | def _iterate_licensed_files(num_lines: int) -> Iterator[Tuple[str, str]]: 100 | source_files = list(Path(SOURCE_DIR).rglob(FILES_WITH_LICENSE_NEEDED_GLOB_EXPR)) 101 | test_files = list(Path(TESTS_DIR).rglob(FILES_WITH_LICENSE_NEEDED_GLOB_EXPR)) 102 | script_files = list(Path(SCRIPTS_DIR).rglob(FILES_WITH_LICENSE_NEEDED_GLOB_EXPR)) 103 | for path in chain(source_files, test_files, script_files): 104 | filename = str(path) 105 | file_head = _get_n_lines_of_file(filename=filename, num_lines=num_lines) 106 | yield filename, file_head 107 | 108 | 109 | def _get_files_list_str(files_list: List[str]) -> str: 110 | return "\n".join([f"\t{filename}" for filename in files_list]) 111 | 112 | 113 | def check_license() -> None: 114 | """Raises an exception if there are any files with no license present at the top of the file""" 115 | contents, num_lines = _get_license_content_any_year() 116 | failed_files = [] 117 | for filename, file_head in _iterate_licensed_files(num_lines=num_lines): 118 | if not any(map(lambda license: file_head.startswith(license), contents)): 119 | failed_files.append(filename) 120 | if failed_files: 121 | print( 122 | "Some files did not have the license header included!\n\n" 123 | + _get_files_list_str(failed_files) 124 | + "\n\nRun `poe license` and commit to fix the issue" 125 | ) 126 | else: 127 | print(f"All {FILES_WITH_LICENSE_NEEDED_GLOB_EXPR} have license header") 128 | sys.exit(len(failed_files)) 129 | 130 | 131 | def _add_license_to_file(filepath: str, license_content: str) -> None: 132 | """Adds license header to top of a file""" 133 | with open(filepath, "r", encoding="utf-8") as old: 134 | unlink(filepath) 135 | with open(filepath, "w", encoding="utf-8") as new: 136 | new.write(license_content + "\n\n\n") 137 | shutil.copyfileobj(old, new) 138 | 139 | 140 | def license() -> None: 141 | """Adds license header to any files that are missing it""" 142 | contents, num_lines = _get_license_content_any_year() 143 | current_year_license = contents[-1] 144 | updated_files = [] 145 | for filename, file_head in _iterate_licensed_files(num_lines=num_lines): 146 | if not any(map(lambda license: file_head.startswith(license), contents)): 147 | _add_license_to_file(filepath=filename, license_content=current_year_license) 148 | updated_files.append(filename) 149 | 150 | if updated_files: 151 | print(f"Added license to the following files:\n\n{_get_files_list_str(updated_files)}") 152 | else: 153 | print(f"All {FILES_WITH_LICENSE_NEEDED_GLOB_EXPR} have license header") 154 | -------------------------------------------------------------------------------- /compute_modules/sources_v2/_sources.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | from datetime import datetime 17 | from functools import cache 18 | from typing import Any, Optional 19 | 20 | from ._api import DEFAULT_CA_BUNDLE, JAVA_OFFSET_DATETIME_FORMAT, SERVICE_DISCOVERY_PATH, MountedSourceConfig 21 | from ._back_compat import get_mounted_sources 22 | 23 | 24 | def __get_on_prem_proxy_service_uris() -> list[str]: 25 | try: 26 | import yaml # type: ignore[import-untyped] 27 | except ImportError: 28 | raise ImportError( 29 | "the sources extras is not installed. Please install it with `pip install foundry-compute-modules[sources]`" 30 | ) 31 | 32 | with open(os.environ[SERVICE_DISCOVERY_PATH], "r") as f: 33 | service_discovery = yaml.safe_load(f) 34 | on_prem_proxy_uris: list[str] = service_discovery.get("on_prem_proxy", []) 35 | return on_prem_proxy_uris 36 | 37 | 38 | @cache 39 | def get_source(source_api_name: str): # type: ignore[no-untyped-def] 40 | try: 41 | from external_systems.sources import ( 42 | AwsCredentials, 43 | ClientCertificate, 44 | GcpOauthCredentials, 45 | HttpsConnectionParameters, 46 | OauthCredentials, 47 | Source, 48 | SourceCredentials, 49 | SourceParameters, 50 | ) 51 | except ImportError: 52 | raise ImportError( 53 | "the sources extras is not installed. Please install it with `pip install foundry-compute-modules[sources]`" 54 | ) 55 | 56 | def convert_resolved_source_credentials( 57 | credentials: Optional[Any], 58 | ) -> Optional[SourceCredentials]: 59 | if credentials is None: 60 | return None 61 | 62 | cloud_credentials = _maybe_get_cloud_credentials(credentials) 63 | if cloud_credentials is not None: 64 | return cloud_credentials 65 | 66 | gcp_oauth_credentials = _maybe_get_gcp_oauth_credentials(credentials) 67 | if gcp_oauth_credentials is not None: 68 | return gcp_oauth_credentials 69 | 70 | oauth2_credentials = _maybe_get_oauth_credentials(credentials) 71 | if oauth2_credentials is not None: 72 | return oauth2_credentials 73 | 74 | return None 75 | 76 | def _maybe_get_oauth_credentials( 77 | credentials: Any, 78 | ) -> Optional[SourceCredentials]: 79 | oauth_credentials = credentials.get("oauth2Credentials") 80 | if oauth_credentials is None: 81 | return None 82 | 83 | return OauthCredentials( 84 | access_token=oauth_credentials.get("accessToken"), 85 | expiration=datetime.strptime(oauth_credentials.get("expiration"), JAVA_OFFSET_DATETIME_FORMAT), 86 | ) 87 | 88 | def _maybe_get_gcp_oauth_credentials( 89 | credentials: Any, 90 | ) -> Optional[SourceCredentials]: 91 | gcp_oauth_credentials = credentials.get("gcpOauthCredentials", None) 92 | if gcp_oauth_credentials is None: 93 | return None 94 | 95 | return GcpOauthCredentials( 96 | access_token=gcp_oauth_credentials.get("accessToken"), 97 | expiration=datetime.strptime(gcp_oauth_credentials.get("expiration"), JAVA_OFFSET_DATETIME_FORMAT), 98 | ) 99 | 100 | def _maybe_get_cloud_credentials( 101 | credentials: Any, 102 | ) -> Optional[SourceCredentials]: 103 | cloud_credentials = credentials.get("cloudCredentials", None) 104 | if cloud_credentials is None: 105 | return None 106 | 107 | aws_credentials = cloud_credentials.get("awsCredentials", None) 108 | if aws_credentials is None: 109 | return None 110 | 111 | session_credentials = aws_credentials.get("sessionCredentials", None) 112 | if session_credentials is not None: 113 | return AwsCredentials( 114 | access_key_id=session_credentials.get("accessKeyId"), 115 | secret_access_key=session_credentials.get("secretAccessKey"), 116 | session_token=session_credentials.get("sessionToken"), 117 | expiration=datetime.strptime(session_credentials.get("expiration"), JAVA_OFFSET_DATETIME_FORMAT), 118 | ) 119 | 120 | basic_credentials = aws_credentials.get("basicCredentials", None) 121 | if basic_credentials is not None: 122 | return AwsCredentials( 123 | access_key_id=basic_credentials.get("accessKeyId"), 124 | secret_access_key=basic_credentials.get("secretAccessKey"), 125 | ) 126 | 127 | return None 128 | 129 | mounted_source: Optional[MountedSourceConfig] = get_mounted_sources().get(source_api_name, None) 130 | if mounted_source is None: 131 | raise ValueError(f"Source {source_api_name} not found") 132 | 133 | client_certificate = ( 134 | ClientCertificate( 135 | pem_certificate=mounted_source.client_certificate.pem_certificate, 136 | pem_private_key=mounted_source.client_certificate.pem_private_key, 137 | ) 138 | if mounted_source.client_certificate is not None 139 | else None 140 | ) 141 | 142 | source_parameters = SourceParameters( 143 | secrets=mounted_source.secrets, 144 | proxy_token=mounted_source.proxy_token, 145 | https_connections=( 146 | { 147 | "http_connection": HttpsConnectionParameters( 148 | url=mounted_source.http_connection_config.url, 149 | headers=mounted_source.http_connection_config.auth_headers, 150 | query_params=mounted_source.http_connection_config.query_parameters, 151 | ) 152 | } 153 | if mounted_source.http_connection_config is not None 154 | else {} 155 | ), 156 | server_certificates=mounted_source.server_certificates, 157 | client_certificate=client_certificate, 158 | resolved_source_credentials=convert_resolved_source_credentials(mounted_source.resolved_credentials), 159 | ) 160 | 161 | on_prem_proxy_service_uris = __get_on_prem_proxy_service_uris() 162 | 163 | return Source( 164 | source_parameters=source_parameters, 165 | source_configuration=mounted_source.source_configuration, 166 | on_prem_proxy_service_uris=on_prem_proxy_service_uris, 167 | egress_proxy_service_uris=[], 168 | egress_proxy_token=None, 169 | ca_bundle_path=os.environ.get(DEFAULT_CA_BUNDLE), 170 | ) 171 | -------------------------------------------------------------------------------- /compute_modules/bin/static_inference/infer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import argparse 16 | import ast 17 | import importlib 18 | import inspect 19 | import json 20 | import logging 21 | import os 22 | import pkgutil 23 | import sys 24 | import types 25 | from typing import Any, Dict, Iterator, List, Optional, Set 26 | 27 | import compute_modules.startup 28 | from compute_modules.bin.ontology._config_path import get_ontology_config_file 29 | from compute_modules.function_registry.function import Function 30 | from compute_modules.function_registry.function_registry import add_function, add_functions 31 | from compute_modules.function_registry.types import ComputeModuleFunctionSchema 32 | 33 | LOGGER = logging.getLogger(__name__) 34 | 35 | 36 | def infer( 37 | src_dir: str, 38 | api_name_type_id_mapping: Dict[str, str], 39 | ) -> List[ComputeModuleFunctionSchema]: 40 | # Disables automatically starting compute module upon importing function annotations 41 | compute_modules.startup.DISABLE_STARTUP = True 42 | 43 | if src_dir not in sys.path: 44 | sys.path.append(src_dir) 45 | 46 | py_modules: Set[types.ModuleType] = set(_import_python_modules(src_dir)) 47 | cm_functions: List[Function] = list(_discover_functions(py_modules)) 48 | _validate_functions(cm_functions) 49 | return _parse_function_schemas(cm_functions, api_name_type_id_mapping) 50 | 51 | 52 | def _import_python_modules(directory: str) -> Iterator[types.ModuleType]: 53 | for module in pkgutil.walk_packages([directory]): 54 | LOGGER.debug(f"Found {repr(module)}") 55 | if not module.ispkg: 56 | LOGGER.debug(f"Importing module {module.name}") 57 | yield importlib.import_module(module.name) 58 | 59 | 60 | def _discover_functions( 61 | py_modules: Set[types.ModuleType], 62 | ) -> Iterator[Function]: 63 | for module in py_modules: 64 | yield from _discover_decorated_functions(module) 65 | yield from _discover_manually_registered_functions(module) 66 | 67 | 68 | def _discover_decorated_functions(py_module: types.ModuleType) -> Iterator[Function]: 69 | module_attrs = [getattr(py_module, attr) for attr in dir(py_module)] 70 | for attr in module_attrs: 71 | if inspect.getmodule(attr) is py_module and isinstance(attr, Function): 72 | LOGGER.debug(f"Located function {attr.__name__} in module {py_module.__name__}") 73 | yield attr 74 | 75 | 76 | def _discover_manually_registered_functions(py_module: types.ModuleType) -> Iterator[Function]: 77 | try: 78 | source = inspect.getsource(py_module) 79 | except OSError: 80 | LOGGER.warning(f"Could not read source for module {py_module.__name__}") 81 | # https://stackoverflow.com/questions/13243766/how-to-define-an-empty-generator-function 82 | return 83 | yield 84 | 85 | syntax_tree = ast.parse(source) 86 | 87 | for node in ast.walk(syntax_tree): 88 | if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name): 89 | continue 90 | # Manual function registration entry points 91 | if not (node.func.id == add_function.__name__ or node.func.id == add_functions.__name__): 92 | continue 93 | LOGGER.debug(f"Found call to {node.func.id} in module {py_module.__name__}") 94 | 95 | for arg in node.args: 96 | if not isinstance(arg, ast.Name): 97 | continue 98 | fn = getattr(py_module, arg.id, None) 99 | if not callable(fn): 100 | continue 101 | 102 | LOGGER.debug(f"Located function {fn.__name__} in module {py_module.__name__}") 103 | # Extracting ontology types if `edits=[...]` was provided 104 | edits_arg = next(filter(lambda k: k.arg == "edits", node.keywords), None) 105 | parsed_edits: Set[Any] = set() 106 | _maybe_add_edits(py_module=py_module, edits_arg=edits_arg, parsed_edits=parsed_edits) 107 | yield Function(fn, list(parsed_edits)) 108 | 109 | 110 | def _maybe_add_edits( 111 | py_module: types.ModuleType, 112 | edits_arg: Optional[Any], 113 | parsed_edits: Set[Any], 114 | ) -> None: 115 | if not edits_arg: 116 | return 117 | # edits=[...] literal list syntax 118 | if isinstance(edits_arg, ast.keyword) and isinstance(edits_arg.value, ast.List): 119 | for edit in edits_arg.value.elts: 120 | if not isinstance(edit, ast.Name): 121 | continue 122 | parsed_edit = getattr(py_module, edit.id, None) 123 | if parsed_edit and hasattr(parsed_edit, "api_name") and callable(parsed_edit.api_name): 124 | parsed_edits.add(parsed_edit) 125 | # edits=SOME_VAR syntax 126 | if isinstance(edits_arg, ast.keyword) and isinstance(edits_arg.value, ast.Name): 127 | edits_list = getattr(py_module, edits_arg.value.id, None) 128 | if edits_list and isinstance(edits_list, list): 129 | for edit_type in edits_list: 130 | if edit_type and hasattr(edit_type, "api_name") and callable(edit_type.api_name): 131 | parsed_edits.add(edit_type) 132 | 133 | 134 | def _validate_functions(functions: List[Function]) -> None: 135 | seen_functions: Set[str] = set() 136 | duplicate_functions: List[str] = [] 137 | 138 | for f in functions: 139 | if f.__name__ in seen_functions: 140 | duplicate_functions.append(f.__name__) 141 | else: 142 | seen_functions.add(f.__name__) 143 | 144 | if len(duplicate_functions) > 0: 145 | raise ValueError(f"Duplicate function(s) found: {duplicate_functions}") 146 | 147 | 148 | def _parse_function_schemas( 149 | functions: List[Function], 150 | api_name_type_id_mapping: Dict[str, str], 151 | ) -> List[ComputeModuleFunctionSchema]: 152 | parsed_schemas = [] 153 | for function in functions: 154 | LOGGER.debug(f"Serialising function {function.__name__}") 155 | parsed_schemas.append(function.get_function_schema(api_name_type_id_mapping)) 156 | return parsed_schemas 157 | 158 | 159 | def main() -> None: 160 | parser = argparse.ArgumentParser() 161 | parser.add_argument( 162 | "source", 163 | help="Path to the source directory of your compute module", 164 | ) 165 | parser.add_argument( 166 | "--ontology-metadata-config", 167 | required=False, 168 | help="Path to a configuration file that is used to determine type information for OSDK types. " 169 | + "Only needed if ontology edits are returned by any functions.", 170 | dest="ontology_metadata_config_file", 171 | default=None, 172 | ) 173 | arguments = parser.parse_args() 174 | config_file_path = get_ontology_config_file(arguments.ontology_metadata_config_file) 175 | api_name_type_id_mapping = _get_api_name_type_id_mapping(config_file_path) 176 | print( 177 | json.dumps( 178 | infer( 179 | src_dir=arguments.source, 180 | api_name_type_id_mapping=api_name_type_id_mapping, 181 | ) 182 | ) 183 | ) 184 | 185 | 186 | def _get_api_name_type_id_mapping(config_file_path: str) -> dict[str, str]: 187 | if not os.path.isfile(config_file_path): 188 | return {} 189 | with open(config_file_path) as f: 190 | config_data = json.load(f) 191 | return config_data.get("apiNameToTypeId", {}) # type: ignore[no-any-return] 192 | 193 | 194 | if __name__ == "__main__": 195 | main() 196 | -------------------------------------------------------------------------------- /tests/function_registry/test_function_schema_parser.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | from typing import Optional 17 | 18 | import pytest 19 | 20 | from compute_modules.function_registry.function_schema_parser import parse_function_schema 21 | from compute_modules.function_registry.types import ComputeModuleFunctionSchema, FunctionOutputType 22 | from tests.function_registry.dummy_app import ( 23 | DummyInput, 24 | ParentClass, 25 | dummy_func_1, 26 | dummy_func_2, 27 | dummy_func_3, 28 | dummy_func_4, 29 | dummy_func_5, 30 | ) 31 | from tests.function_registry.dummy_app_with_issues import ( 32 | dummy_args_init, 33 | dummy_kwargs_init, 34 | dummy_no_init_hints, 35 | dummy_no_type_hints, 36 | ) 37 | 38 | EXPECTED_OUTPUT_1 = { 39 | "single": { 40 | "dataType": { 41 | "anonymousCustomType": { 42 | "fields": { 43 | "res1": {"boolean": {}, "type": "boolean"}, 44 | "res2": { 45 | "map": { 46 | "keysType": {"string": {}, "type": "string"}, 47 | "valuesType": {"float": {}, "type": "float"}, 48 | }, 49 | "type": "map", 50 | }, 51 | "res3": {"list": {"elementsType": {"timestamp": {}, "type": "timestamp"}}, "type": "list"}, 52 | } 53 | }, 54 | "type": "anonymousCustomType", 55 | } 56 | }, 57 | "type": "single", 58 | } 59 | 60 | EXPECTED_OUTPUT_3 = { 61 | "single": { 62 | "dataType": { 63 | "integer": {}, 64 | "type": "integer", 65 | } 66 | }, 67 | "type": "single", 68 | } 69 | 70 | EXPECTED_OUTPUT_4 = { 71 | "single": { 72 | "dataType": { 73 | "list": { 74 | "elementsType": {"string": {}, "type": "string"}, 75 | }, 76 | "type": "list", 77 | } 78 | }, 79 | "type": "single", 80 | } 81 | 82 | EXPECTED_INPUTS = [ 83 | { 84 | "name": "parent_class", 85 | "dataType": { 86 | "anonymousCustomType": { 87 | "fields": { 88 | "some_flag": {"boolean": {}, "type": "boolean"}, 89 | "some_value": {"integer": {}, "type": "integer"}, 90 | "child": { 91 | "anonymousCustomType": { 92 | "fields": { 93 | "timestamp": {"timestamp": {}, "type": "timestamp"}, 94 | "some_value": {"float": {}, "type": "float"}, 95 | "another_optional_field": { 96 | "optionalType": {"wrappedType": {"string": {}, "type": "string"}}, 97 | "type": "optionalType", 98 | }, 99 | } 100 | }, 101 | "type": "anonymousCustomType", 102 | }, 103 | } 104 | }, 105 | "type": "anonymousCustomType", 106 | }, 107 | "required": True, 108 | "constraints": [], 109 | }, 110 | { 111 | "name": "optional_field", 112 | "dataType": {"optionalType": {"wrappedType": {"string": {}, "type": "string"}}, "type": "optionalType"}, 113 | "required": True, 114 | "constraints": [], 115 | }, 116 | { 117 | "name": "set_field", 118 | "dataType": {"set": {"elementsType": {"date": {}, "type": "date"}}, "type": "set"}, 119 | "required": True, 120 | "constraints": [], 121 | }, 122 | { 123 | "name": "map_field", 124 | "dataType": { 125 | "map": {"keysType": {"binary": {}, "type": "binary"}, "valuesType": {"decimal": {}, "type": "decimal"}}, 126 | "type": "map", 127 | }, 128 | "required": True, 129 | "constraints": [], 130 | }, 131 | { 132 | "name": "datetime_list", 133 | "dataType": {"list": {"elementsType": {"timestamp": {}, "type": "timestamp"}}, "type": "list"}, 134 | "required": True, 135 | "constraints": [], 136 | }, 137 | {"name": "some_flag", "dataType": {"boolean": {}, "type": "boolean"}, "required": True, "constraints": []}, 138 | { 139 | "name": "optional_default_value_field", 140 | "dataType": {"optionalType": {"wrappedType": {"string": {}, "type": "string"}}, "type": "optionalType"}, 141 | "required": True, 142 | "constraints": [], 143 | }, 144 | ] 145 | 146 | 147 | def test_function_schema_parser() -> None: 148 | """Test the happy path for parse_function_schemas_from_module""" 149 | parse_result = parse_function_schema(dummy_func_1, "dummy_func_1", [], {}) 150 | assert parse_result.function_schema["functionName"] == "dummy_func_1" 151 | assert parse_result.function_schema["output"] == EXPECTED_OUTPUT_1 152 | assert len(parse_result.function_schema["inputs"]) == len(EXPECTED_INPUTS) 153 | assert parse_result.class_node is not None 154 | assert parse_result.class_node["constructor"] is DummyInput 155 | assert parse_result.function_schema["inputs"] == EXPECTED_INPUTS 156 | assert parse_result.class_node["children"] is not None 157 | assert parse_result.class_node["children"]["parent_class"]["constructor"] is ParentClass 158 | assert parse_result.class_node["children"]["optional_field"]["constructor"] is Optional 159 | assert parse_result.class_node["children"]["set_field"]["constructor"] is set 160 | assert parse_result.class_node["children"]["map_field"]["constructor"] is dict 161 | assert parse_result.class_node["children"]["some_flag"]["constructor"] is bool 162 | assert parse_result.class_node["children"]["some_flag"]["children"] is None 163 | assert parse_result.class_node["children"]["optional_default_value_field"]["constructor"] is Optional 164 | assert parse_result.is_context_typed is False 165 | 166 | 167 | def test_function_schema_parser_no_type_hints() -> None: 168 | """Test 'happy' path, but on a function with no type hints""" 169 | parse_result = parse_function_schema(dummy_func_2, "dummy_func_2", [], {}) 170 | assert parse_result.class_node is None 171 | assert parse_result.is_context_typed is False 172 | assert parse_result.function_schema == ComputeModuleFunctionSchema( 173 | functionName="dummy_func_2", 174 | inputs=[], 175 | output=FunctionOutputType( 176 | type="single", 177 | single={ 178 | "dataType": {"type": "string", "string": {}}, 179 | }, 180 | ), 181 | ontologyProvenance=None, 182 | ) 183 | 184 | 185 | def test_function_schema_parser_context_output_only() -> None: 186 | """Test 'happy' path for a function that uses type hints only for the context & return type""" 187 | parse_result = parse_function_schema(dummy_func_3, "dummy_func_3", [], {}) 188 | assert parse_result.class_node is None 189 | assert parse_result.is_context_typed 190 | assert parse_result.function_schema["functionName"] == "dummy_func_3" 191 | assert parse_result.function_schema["inputs"] == [] 192 | assert parse_result.function_schema["output"] == EXPECTED_OUTPUT_3 193 | 194 | 195 | def test_function_schema_parser_dict_witout_params() -> None: 196 | """Test 'happy' path for a function that uses type hints only for the context & return type""" 197 | with pytest.raises(ValueError) as exc_info: 198 | parse_function_schema(dummy_func_4, "dummy_func_4", [], {}) 199 | assert "dict type hints must have type parameters provided" in str(exc_info.value) 200 | 201 | 202 | def test_exception_no_type_hints() -> None: 203 | """CM function params should not have classes without type hints""" 204 | with pytest.raises(ValueError) as exc_info: 205 | parse_function_schema(dummy_no_type_hints, "dummy_no_type_hints", [], {}) 206 | assert "type_hints set() must match init args" in str(exc_info.value) 207 | 208 | 209 | def test_exception_no_init_hints() -> None: 210 | """CM function params should not have constructors without type hints""" 211 | with pytest.raises(ValueError) as exc_info: 212 | parse_function_schema(dummy_no_init_hints, "dummy_no_init_hints", [], {}) 213 | assert "Custom Type BadClassNoInitHints should have init args type annotations" in str(exc_info.value) 214 | 215 | 216 | def test_exception_args_init() -> None: 217 | """CM function params should not have constructors that use the `args` keyword""" 218 | with pytest.raises(ValueError) as exc_info: 219 | parse_function_schema(dummy_args_init, "dummy_args_init", [], {}) 220 | assert "The __init__ method should not use *args" in str(exc_info.value) 221 | 222 | 223 | def test_exception_kwargs_init() -> None: 224 | """CM function params should not have constructors that use the `kwargs` keyword""" 225 | with pytest.raises(ValueError) as exc_info: 226 | parse_function_schema(dummy_kwargs_init, "dummy_kwargs_init", [], {}) 227 | assert "The __init__ method should not use **kwargs" in str(exc_info.value) 228 | 229 | 230 | def test_function_schema_parser_generator_output() -> None: 231 | """Test 'happy' path for a function that uses type hints for generator return type""" 232 | parse_result = parse_function_schema(dummy_func_5, "dummy_func_5", [], {}) 233 | assert parse_result.function_schema["functionName"] == "dummy_func_5" 234 | assert parse_result.function_schema["output"] == EXPECTED_OUTPUT_4 235 | -------------------------------------------------------------------------------- /tests/infer/test_infer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 Palantir Technologies, Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | import os 17 | from typing import List 18 | 19 | import pytest 20 | 21 | from compute_modules.bin.static_inference.infer import infer 22 | from compute_modules.function_registry.types import ComputeModuleFunctionSchema 23 | 24 | API_NAME_TYPE_ID_MAPPING = {"DummyOntologyType": "dummy-ontology"} 25 | CURRENT_FILE_PATH = os.path.abspath(__file__) 26 | CURRENT_DIR = os.path.dirname(CURRENT_FILE_PATH) 27 | DECORATED_DIR = os.path.join(CURRENT_DIR, "test_project", "decorated") 28 | MANUALLY_REGISTERED_DIR = os.path.join(CURRENT_DIR, "test_project", "manually_registered") 29 | 30 | 31 | EXPECTED_DECORATED_RES = [ 32 | { 33 | "functionName": "ontology_add_function", 34 | "inputs": [], 35 | "output": { 36 | "type": "single", 37 | "single": { 38 | "dataType": {"type": "list", "list": {"elementsType": {"ontologyEdit": {}, "type": "ontologyEdit"}}} 39 | }, 40 | }, 41 | "ontologyProvenance": {"editedObjects": {"dummy-ontology": {}}, "editedLinks": {}}, 42 | }, 43 | { 44 | "functionName": "return_dict", 45 | "inputs": [], 46 | "output": { 47 | "type": "single", 48 | "single": { 49 | "dataType": { 50 | "type": "map", 51 | "map": { 52 | "keysType": {"type": "string", "string": {}}, 53 | "valuesType": {"type": "string", "string": {}}, 54 | }, 55 | } 56 | }, 57 | }, 58 | "ontologyProvenance": None, 59 | }, 60 | { 61 | "functionName": "return_list", 62 | "inputs": [], 63 | "output": { 64 | "type": "single", 65 | "single": {"dataType": {"type": "list", "list": {"elementsType": {"type": "string", "string": {}}}}}, 66 | }, 67 | "ontologyProvenance": None, 68 | }, 69 | { 70 | "functionName": "return_set", 71 | "inputs": [], 72 | "output": { 73 | "type": "single", 74 | "single": {"dataType": {"type": "set", "set": {"elementsType": {"type": "string", "string": {}}}}}, 75 | }, 76 | "ontologyProvenance": None, 77 | }, 78 | { 79 | "functionName": "return_byte", 80 | "inputs": [{"name": "value", "required": True, "constraints": [], "dataType": {"type": "byte", "byte": {}}}], 81 | "output": {"type": "single", "single": {"dataType": {"type": "byte", "byte": {}}}}, 82 | "ontologyProvenance": None, 83 | }, 84 | { 85 | "functionName": "return_double", 86 | "inputs": [ 87 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "double", "double": {}}} 88 | ], 89 | "output": {"type": "single", "single": {"dataType": {"type": "double", "double": {}}}}, 90 | "ontologyProvenance": None, 91 | }, 92 | { 93 | "functionName": "return_long", 94 | "inputs": [{"name": "value", "required": True, "constraints": [], "dataType": {"type": "long", "long": {}}}], 95 | "output": {"type": "single", "single": {"dataType": {"type": "long", "long": {}}}}, 96 | "ontologyProvenance": None, 97 | }, 98 | { 99 | "functionName": "return_short", 100 | "inputs": [{"name": "value", "required": True, "constraints": [], "dataType": {"type": "short", "short": {}}}], 101 | "output": {"type": "single", "single": {"dataType": {"type": "short", "short": {}}}}, 102 | "ontologyProvenance": None, 103 | }, 104 | { 105 | "functionName": "return_bool", 106 | "inputs": [ 107 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "boolean", "boolean": {}}} 108 | ], 109 | "output": {"type": "single", "single": {"dataType": {"type": "boolean", "boolean": {}}}}, 110 | "ontologyProvenance": None, 111 | }, 112 | { 113 | "functionName": "return_bytes", 114 | "inputs": [ 115 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "binary", "binary": {}}} 116 | ], 117 | "output": {"type": "single", "single": {"dataType": {"type": "binary", "binary": {}}}}, 118 | "ontologyProvenance": None, 119 | }, 120 | { 121 | "functionName": "return_date", 122 | "inputs": [{"name": "value", "required": True, "constraints": [], "dataType": {"type": "date", "date": {}}}], 123 | "output": {"type": "single", "single": {"dataType": {"type": "date", "date": {}}}}, 124 | "ontologyProvenance": None, 125 | }, 126 | { 127 | "functionName": "return_datetime", 128 | "inputs": [ 129 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "timestamp", "timestamp": {}}} 130 | ], 131 | "output": {"type": "single", "single": {"dataType": {"type": "timestamp", "timestamp": {}}}}, 132 | "ontologyProvenance": None, 133 | }, 134 | { 135 | "functionName": "return_decimal", 136 | "inputs": [ 137 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "decimal", "decimal": {}}} 138 | ], 139 | "output": {"type": "single", "single": {"dataType": {"type": "decimal", "decimal": {}}}}, 140 | "ontologyProvenance": None, 141 | }, 142 | { 143 | "functionName": "return_float", 144 | "inputs": [{"name": "value", "required": True, "constraints": [], "dataType": {"type": "float", "float": {}}}], 145 | "output": {"type": "single", "single": {"dataType": {"type": "float", "float": {}}}}, 146 | "ontologyProvenance": None, 147 | }, 148 | { 149 | "functionName": "return_int", 150 | "inputs": [ 151 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "integer", "integer": {}}} 152 | ], 153 | "output": {"type": "single", "single": {"dataType": {"type": "integer", "integer": {}}}}, 154 | "ontologyProvenance": None, 155 | }, 156 | { 157 | "functionName": "return_str", 158 | "inputs": [ 159 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "string", "string": {}}} 160 | ], 161 | "output": {"type": "single", "single": {"dataType": {"type": "string", "string": {}}}}, 162 | "ontologyProvenance": None, 163 | }, 164 | ] 165 | 166 | EXPECTED_MANUALLY_REGISTERED_RES = [ 167 | { 168 | "functionName": "return_dict_length_in_main", 169 | "inputs": [ 170 | { 171 | "name": "value", 172 | "required": True, 173 | "constraints": [], 174 | "dataType": { 175 | "type": "map", 176 | "map": { 177 | "keysType": {"type": "string", "string": {}}, 178 | "valuesType": {"type": "string", "string": {}}, 179 | }, 180 | }, 181 | } 182 | ], 183 | "output": {"type": "single", "single": {"dataType": {"type": "integer", "integer": {}}}}, 184 | "ontologyProvenance": None, 185 | }, 186 | { 187 | "functionName": "return_complex_in_main", 188 | "inputs": [ 189 | { 190 | "name": "points", 191 | "required": True, 192 | "constraints": [], 193 | "dataType": { 194 | "type": "list", 195 | "list": { 196 | "elementsType": { 197 | "type": "anonymousCustomType", 198 | "anonymousCustomType": { 199 | "fields": { 200 | "x": {"type": "float", "float": {}}, 201 | "y": {"type": "float", "float": {}}, 202 | "z": {"type": "float", "float": {}}, 203 | } 204 | }, 205 | } 206 | }, 207 | }, 208 | }, 209 | { 210 | "name": "messages", 211 | "required": True, 212 | "constraints": [], 213 | "dataType": { 214 | "type": "anonymousCustomType", 215 | "anonymousCustomType": { 216 | "fields": { 217 | "messages": { 218 | "type": "list", 219 | "list": { 220 | "elementsType": { 221 | "type": "anonymousCustomType", 222 | "anonymousCustomType": { 223 | "fields": { 224 | "message": {"type": "string", "string": {}}, 225 | "from_id": {"type": "integer", "integer": {}}, 226 | "to_id": {"type": "integer", "integer": {}}, 227 | "timestamp": {"type": "timestamp", "timestamp": {}}, 228 | } 229 | }, 230 | } 231 | }, 232 | } 233 | } 234 | }, 235 | }, 236 | }, 237 | ], 238 | "output": {"type": "single", "single": {"dataType": {"type": "string", "string": {}}}}, 239 | "ontologyProvenance": None, 240 | }, 241 | { 242 | "functionName": "return_number_in_main", 243 | "inputs": [ 244 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "integer", "integer": {}}} 245 | ], 246 | "output": { 247 | "type": "single", 248 | "single": { 249 | "dataType": { 250 | "type": "anonymousCustomType", 251 | "anonymousCustomType": {"fields": {"value": {"type": "integer", "integer": {}}}}, 252 | } 253 | }, 254 | }, 255 | "ontologyProvenance": None, 256 | }, 257 | { 258 | "functionName": "ontology_edit_function", 259 | "inputs": [{"name": "name", "required": True, "constraints": [], "dataType": {"type": "string", "string": {}}}], 260 | "output": { 261 | "type": "single", 262 | "single": { 263 | "dataType": {"type": "list", "list": {"elementsType": {"ontologyEdit": {}, "type": "ontologyEdit"}}} 264 | }, 265 | }, 266 | "ontologyProvenance": {"editedObjects": {"dummy-ontology": {}}, "editedLinks": {}}, 267 | }, 268 | { 269 | "functionName": "return_integer_in_main", 270 | "inputs": [ 271 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "integer", "integer": {}}} 272 | ], 273 | "output": {"type": "single", "single": {"dataType": {"type": "integer", "integer": {}}}}, 274 | "ontologyProvenance": None, 275 | }, 276 | { 277 | "functionName": "mixed_1", 278 | "inputs": [ 279 | {"name": "value", "required": True, "constraints": [], "dataType": {"type": "integer", "integer": {}}} 280 | ], 281 | "output": {"type": "single", "single": {"dataType": {"type": "integer", "integer": {}}}}, 282 | "ontologyProvenance": None, 283 | }, 284 | { 285 | "functionName": "mixed_2", 286 | "inputs": [ 287 | { 288 | "name": "value", 289 | "required": True, 290 | "constraints": [], 291 | "dataType": {"type": "set", "set": {"elementsType": {"type": "integer", "integer": {}}}}, 292 | } 293 | ], 294 | "output": { 295 | "type": "single", 296 | "single": {"dataType": {"type": "set", "set": {"elementsType": {"type": "integer", "integer": {}}}}}, 297 | }, 298 | "ontologyProvenance": None, 299 | }, 300 | { 301 | "functionName": "mixed_3", 302 | "inputs": [{"name": "value", "required": True, "constraints": [], "dataType": {"type": "byte", "byte": {}}}], 303 | "output": {"type": "single", "single": {"dataType": {"type": "byte", "byte": {}}}}, 304 | "ontologyProvenance": None, 305 | }, 306 | { 307 | "functionName": "mixed_4", 308 | "inputs": [{"name": "name", "required": True, "constraints": [], "dataType": {"type": "string", "string": {}}}], 309 | "output": {"type": "single", "single": {"dataType": {"type": "string", "string": {}}}}, 310 | "ontologyProvenance": None, 311 | }, 312 | ] 313 | test_cases = [ 314 | ( 315 | DECORATED_DIR, 316 | EXPECTED_DECORATED_RES, 317 | ), 318 | ( 319 | MANUALLY_REGISTERED_DIR, 320 | EXPECTED_MANUALLY_REGISTERED_RES, 321 | ), 322 | ] 323 | 324 | 325 | @pytest.mark.parametrize("src_dir, expected_result", test_cases) 326 | def test_infer(src_dir: str, expected_result: List[ComputeModuleFunctionSchema]) -> None: 327 | res = infer( 328 | src_dir=src_dir, 329 | api_name_type_id_mapping=API_NAME_TYPE_ID_MAPPING, 330 | ) 331 | assert sorted(res, key=lambda x: x["functionName"]) == sorted(expected_result, key=lambda x: x["functionName"]) 332 | --------------------------------------------------------------------------------