diff --git a/.devcontainer/README.md b/.devcontainer/README.md new file mode 100644 index 000000000..cac7aa57d --- /dev/null +++ b/.devcontainer/README.md @@ -0,0 +1,48 @@ +# WranglesPY development container + +The shared `devcontainer.json` is intentionally host-neutral. VS Code on Windows +and GitHub Codespaces both use the same Python 3.13 container definition. + +## Windows with Docker Desktop + +Docker Desktop's Linux engine must be ready before VS Code evaluates the +devcontainer. From the repository root, run: + +```powershell +.\.devcontainer\wait-for-docker.ps1 +``` + +The helper starts Docker Desktop when necessary and waits until the selected +Docker context can reach a Linux server. When it reports that Docker is ready, +open the repository in VS Code and run: + +```text +Dev Containers: Reopen in Container +``` + +Use `Dev Containers: Rebuild Container` after changing `devcontainer.json` or +the dependency files. + +`Dev Containers: Attach to Running Container...` remains useful as a recovery +option when this repository's container is already running. It is not the +normal workflow because attaching does not necessarily apply this repository's +creation commands and VS Code customizations. + +## GitHub Codespaces + +Create or reopen the Codespace normally. Codespaces reads +`.devcontainer/devcontainer.json` directly. The Windows readiness helper is not +referenced by the configuration and does not run in Codespaces. + +## Quick diagnosis + +Before reopening locally, both commands should succeed: + +```powershell +docker context show +docker version +``` + +`docker version` must show a Linux server section. A missing +`dockerDesktopLinuxEngine` named pipe means Docker Desktop is still starting, +not that `devcontainer.json` selected the wrong engine. diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 62ede1354..707c97fc9 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -1,10 +1,10 @@ // For format details, see https://aka.ms/devcontainer.json. For config options, see the // README at: https://github.com/devcontainers/templates/tree/main/src/python { - "name": "Python 3", + "name": "WranglesPY (Python 3.13)", // Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile - "image": "mcr.microsoft.com/devcontainers/python:1-3.12-bookworm", + "image": "mcr.microsoft.com/devcontainers/python:1-3.13-bookworm", // Features to add to the dev container. More info: https://containers.dev/features. // "features": {}, @@ -12,27 +12,37 @@ // Use 'forwardPorts' to make a list of ports inside the container available locally. // "forwardPorts": [], - "onCreateCommand": "pip install pytest==9.0.2 lorem pytest-mock", - "postCreateCommand": "pip3 install --user -r requirements.txt", + "remoteUser": "vscode", + "remoteEnv": { + "PATH": "${containerEnv:PATH}:/home/vscode/.local/bin" + }, + + // Mirror the full CI test environment and install this checkout in editable mode. + "postCreateCommand": "python -m pip install --user --upgrade pip && python -m pip install --user -r requirements-full.txt pytest==9.0.2 pytest-mock -e .", // Configure tool-specific properties. "customizations": { "vscode": { "extensions": [ "redhat.vscode-yaml", - "GitHub.copilot" + "GitHub.copilot", + "editorconfig.editorconfig", + "ms-python.python", + "ms-python.vscode-pylance" ], "settings": { + "files.eol": "\n", + "files.associations": { + "*.recipe": "yaml" + }, "yaml.schemas": { "https://public.wrangle.works/schema/recipes/schema.json": ["*.wrgl.yml", "*.wrgl.yaml", "*.recipe"] }, + "python.defaultInterpreterPath": "/usr/local/bin/python", "python.testing.pytestArgs": ["tests"], "python.testing.unittestEnabled": false, - "python.testing.pytestEnabled": true + "python.testing.pytestEnabled": true } } } - - // Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root. - // "remoteUser": "root" } diff --git a/.devcontainer/wait-for-docker.ps1 b/.devcontainer/wait-for-docker.ps1 new file mode 100644 index 000000000..d30ef0108 --- /dev/null +++ b/.devcontainer/wait-for-docker.ps1 @@ -0,0 +1,52 @@ +[CmdletBinding()] +param( + [ValidateRange(10, 600)] + [int]$TimeoutSeconds = 180 +) + +$ErrorActionPreference = 'Stop' + +if (-not (Get-Command docker -ErrorAction SilentlyContinue)) { + throw 'Docker CLI was not found. Install or repair Docker Desktop before opening the dev container.' +} + +function Get-DockerServerOs { + $serverOs = docker version --format '{{.Server.Os}}' 2>$null + if ($LASTEXITCODE -eq 0) { + return $serverOs.Trim() + } + + return $null +} + +$serverOs = Get-DockerServerOs +if (-not $serverOs) { + $dockerDesktopPath = Join-Path $env:ProgramFiles 'Docker\Docker\Docker Desktop.exe' + if (-not (Test-Path -LiteralPath $dockerDesktopPath)) { + throw "Docker Desktop is not responding and was not found at '$dockerDesktopPath'." + } + + Write-Host 'Starting Docker Desktop...' + Start-Process -FilePath $dockerDesktopPath +} + +$deadline = (Get-Date).AddSeconds($TimeoutSeconds) +Write-Host "Waiting up to $TimeoutSeconds seconds for the Docker Linux engine..." + +do { + $serverOs = Get-DockerServerOs + if ($serverOs -eq 'linux') { + $context = docker context show + Write-Host "Docker is ready (context: $context, server: linux)." + Write-Host "In VS Code, run 'Dev Containers: Reopen in Container'." + exit 0 + } + + if ($serverOs) { + throw "Docker is running a '$serverOs' engine. Switch Docker Desktop to Linux containers and try again." + } + + Start-Sleep -Seconds 2 +} while ((Get-Date) -lt $deadline) + +throw "Docker Desktop did not expose its Linux engine within $TimeoutSeconds seconds. Check Docker Desktop before retrying." diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 000000000..e71865200 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,13 @@ +root = true + +[*] +charset = utf-8 +end_of_line = lf +insert_final_newline = true +trim_trailing_whitespace = true + +[*.{bat,cmd}] +end_of_line = crlf + +[*.md] +trim_trailing_whitespace = false \ No newline at end of file diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 000000000..d324c55e1 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,19 @@ +# Keep all text files LF in both Git and the working tree +* text=auto eol=lf + +# Windows command scripts genuinely require/prefer CRLF +*.bat text eol=crlf +*.cmd text eol=crlf + +# Explicit binary exclusions +*.png -text +*.jpg -text +*.jpeg -text +*.gif -text +*.ico -text +*.pdf -text +*.zip -text +*.tar -text +*.gz -text +*.bz2 -text +*.7z -text \ No newline at end of file diff --git a/.github/agents/test.agent.md b/.github/agents/test.agent.md index 9995c1a58..a08c1ac2b 100644 --- a/.github/agents/test.agent.md +++ b/.github/agents/test.agent.md @@ -27,7 +27,7 @@ You are a **Software Development Engineer in Test (SDET)** for a Python data wra ## šŸ“š Project Knowledge ### Tech Stack -- **Core:** Python 3.10–3.13, pytest, pandas, numpy, polars +- **Core:** Python 3.11–3.13, pytest, pandas, numpy, polars - **Data:** SQLAlchemy, boto3, pymongo, Pydantic - **Templates:** Jinja2 diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index e46432932..4cd718307 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -13,14 +13,14 @@ ## Tech Stack -- **Python:** 3.10, 3.11, 3.12, 3.13 (multi-version support) +- **Python:** 3.11, 3.12, 3.13 (multi-version support) - **Core Dependencies:** pandas (>=2.0), numpy, polars (1.33.0), pyyaml - **Database Connectors:** sqlalchemy, pymssql, psycopg2-binary, pymysql, pymongo - **Cloud/External:** boto3 (AWS S3), simple-salesforce, fabric (SFTP) - **Data Formats:** openpyxl (Excel), xlsxwriter - **AI/ML:** OpenAI integration, Hugging Face models -- **Testing:** pytest (7.4.4), pytest-mock, lorem (test data generation) -- **Containerization:** Docker (Python 3.10.16-slim-bookworm base) +- **Testing:** pytest (9.0.2), pytest-mock, lorem (test data generation) +- **Containerization:** Production Docker image uses Python 3.11-slim-bookworm; development container uses Python 3.13-bookworm ## Project Structure @@ -57,11 +57,12 @@ WranglesPY/ ## Installation & Setup -### Standard Installation +### Development Installation ```bash pip install --upgrade pip -pip install pytest==7.4.4 lorem pytest-mock -pip install -r requirements.txt +pip install pytest==9.0.2 pytest-mock +pip install -r requirements-full.txt +pip install -e . ``` ### macOS-specific Requirements @@ -74,9 +75,9 @@ pip install -r requirements.txt ### Development Container The project includes a `.devcontainer/devcontainer.json` for VS Code: -- Base image: `mcr.microsoft.com/devcontainers/python:1-3.12-bullseye` -- Auto-installs pytest, lorem, pytest-mock on creation -- Includes YAML schema validation for `.wrgl.yml` files +- Base image: `mcr.microsoft.com/devcontainers/python:1-3.13-bookworm` +- Auto-installs the full test dependencies and the package in editable mode +- Includes YAML schema validation for `.wrgl.yml` and `.recipe` files - Configured for pytest test discovery ## Testing @@ -173,7 +174,7 @@ def test_function_error(): ## Recipe System ### Recipe File Format -Recipes use YAML with `.wrgl.yml` or `.wrgl.yaml` extensions: +Recipes use YAML with `.wrgl.yml`, `.wrgl.yaml` or `.recipe` extensions: ```yaml read: @@ -197,10 +198,10 @@ wrangles.recipe recipe.wrgl.yml # From Python import wrangles -wrangles.recipe.run('recipe.wrgl.yml') +wrangles.recipe.run('my_recipe.wrgl.yml') # With custom functions -wrangles.recipe recipe.wrgl.yml -f custom_functions.py +wrangles.recipe my_other_recipe.recipe -f custom_functions.py ``` ### Custom Functions @@ -213,8 +214,8 @@ Custom functions can be added to recipes: ### GitHub Actions Workflows - **publish-main.yml:** Main CI pipeline - - Pytest on multiple OS (Ubuntu, Windows, macOS-14, macOS-latest) - - Tests Python 3.10, 3.11, 3.12, 3.13 + - Pytest on multiple OS (Ubuntu, Windows) + - Tests Python 3.11, 3.13 - Test pip installation - Generate and test JSON schema - Build and push Docker image diff --git a/.github/workflows/publish-dev-rc.yml b/.github/workflows/publish-dev-rc.yml index b1d29a055..97c0393e9 100644 --- a/.github/workflows/publish-dev-rc.yml +++ b/.github/workflows/publish-dev-rc.yml @@ -129,67 +129,67 @@ jobs: pip install . wrangles.recipe tests/samples/generate-data.wrgl.yml - # # ── 3. Generate and test schema ────────────────────────────────────────── - # test-generate-schema: - # name: Generate and Test Schema - # runs-on: ubuntu-latest - # needs: [compute-version] - # permissions: - # contents: read - # steps: - # - name: Checkout Repository - # uses: actions/checkout@v5 - - # - name: Set up Python - # uses: actions/setup-python@v6 - # with: - # python-version: '3.12' - - # - name: Install Dependencies - # run: | - # pip install -r requirements-full.txt - # pip install jsonschema - - # - name: Generate and Test Schema - # run: cd schema && python generate_recipe_schema.py - - # - name: Save schema as schema_dev.json - # run: cp schema/schema.json schema/schema_dev.json - - # - name: Checkout wrangleworks.github.io - # uses: actions/checkout@v5 - # with: - # repository: wrangleworks/wrangleworks.github.io - # token: ${{ secrets.CROSS_REPO_PAT_V2 }} - # path: wrangleworks.github.io - - # - name: Copy schema_dev.json to wrangleworks.github.io - # run: cp schema/schema_dev.json wrangleworks.github.io/schema/recipes/schema_dev.json - - # - name: Commit and push schema_dev.json - # working-directory: wrangleworks.github.io - # run: | - # git config user.name "mborodii-prog" - # git config user.email "mborodii@binariks.com" - # git add schema/recipes/schema_dev.json - # if git diff --cached --quiet; then - # echo "No changes to schema_dev.json, skipping commit." - # else - # git commit -m "Update dev recipe schema" - # git push - # fi - - # - name: Save Schema as Artifact - # uses: actions/upload-artifact@v6 - # with: - # name: schema - # path: schema/schema_dev.json + # ── 3. Generate and test schema ────────────────────────────────────────── + test-generate-schema: + name: Generate and Test Schema + runs-on: ubuntu-latest + needs: [compute-version] + permissions: + contents: read + steps: + - name: Checkout Repository + uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.12' + + - name: Install Dependencies + run: | + pip install -r requirements-full.txt + pip install jsonschema + + - name: Generate and Test Schema + run: cd schema && python generate_recipe_schema.py + + - name: Save schema as schema_dev.json + run: cp schema/schema.json schema/schema_dev.json + + - name: Checkout wrangleworks.github.io + uses: actions/checkout@v5 + with: + repository: wrangleworks/wrangleworks.github.io + token: ${{ secrets.CROSS_REPO_PAT_V2 }} + path: wrangleworks.github.io + + - name: Copy schema_dev.json to wrangleworks.github.io + run: cp schema/schema_dev.json wrangleworks.github.io/schema/recipes/schema_dev.json + + - name: Commit and push schema_dev.json + working-directory: wrangleworks.github.io + run: | + git config user.name "mborodii-prog" + git config user.email "mborodii@binariks.com" + git add schema/recipes/schema_dev.json + if git diff --cached --quiet; then + echo "No changes to schema_dev.json, skipping commit." + else + git commit -m "Update dev recipe schema" + git push + fi + + - name: Save Schema as Artifact + uses: actions/upload-artifact@v6 + with: + name: schema + path: schema/schema_dev.json # ── 4. Publish Python package to CodeArtifact ──────────────────────────── publish-codeartifact: name: Publish RC Package runs-on: ubuntu-latest - needs: [compute-version, test-pip-install] + needs: [compute-version, test-pip-install, test-generate-schema] permissions: contents: read id-token: write # required for OIDC AssumeRoleWithWebIdentity diff --git a/.github/workflows/publish-dev.yml b/.github/workflows/publish-dev.yml index 38c59ca9d..4152a40ad 100644 --- a/.github/workflows/publish-dev.yml +++ b/.github/workflows/publish-dev.yml @@ -18,19 +18,18 @@ jobs: contents: read steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - name: Install Dependencies run: | python -m pip install --upgrade pip pip install pytest==9.0.2 lorem pytest-mock - pip install google-genai - pip install -r requirements.txt + pip install -r requirements-full.txt pip install . - name: Run Tests @@ -51,12 +50,12 @@ jobs: contents: read steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - name: Test pip install run: | @@ -72,13 +71,13 @@ jobs: packages: write steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 # Login against a Docker registry except on PR # https://github.com/docker/login-action - name: Log into registry ${{ env.REGISTRY }} if: github.event_name != 'pull_request' - uses: docker/login-action@v3 + uses: docker/login-action@v4 with: registry: ${{ env.REGISTRY }} username: ${{ github.actor }} @@ -88,14 +87,14 @@ jobs: # https://github.com/docker/metadata-action - name: Extract Docker metadata id: meta - uses: docker/metadata-action@v5 + uses: docker/metadata-action@v6 with: images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} # Build and push Docker image with Buildx (don't push on PR) # https://github.com/docker/build-push-action - name: Build and push Docker image - uses: docker/build-push-action@v5 + uses: docker/build-push-action@v7 with: context: . push: ${{ github.event_name != 'pull_request' }} @@ -119,7 +118,7 @@ jobs: GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Remove wrangles folder run: rm -r wrangles diff --git a/.github/workflows/publish-main.yml b/.github/workflows/publish-main.yml index 1dd088902..d1c929842 100644 --- a/.github/workflows/publish-main.yml +++ b/.github/workflows/publish-main.yml @@ -16,16 +16,16 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - os: [ubuntu-latest, windows-latest, macos-14] - python-version: ['3.10', '3.11', '3.12', '3.13'] + os: [ubuntu-latest, windows-latest] + python-version: ['3.11', '3.13'] permissions: contents: read steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} @@ -33,49 +33,9 @@ jobs: run: | python -m pip install --upgrade pip pip install pytest==9.0.2 lorem pytest-mock - pip install google-genai - pip install -r requirements.txt + pip install -r requirements-full.txt pip install . - - - name: Run Tests - run: pytest - env: - WRANGLES_USER: ${{ secrets.WRANGLES_USER }} - WRANGLES_PASSWORD: ${{ secrets.WRANGLES_PASSWORD }} - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} - HUGGINGFACE_TOKEN: ${{ secrets.HUGGINGFACE_TOKEN }} - OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} - SERPAPI_API_KEY: ${{ secrets.SERPAPI_API_KEY }} - GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} - - pytest-macos-latest: - runs-on: ${{ matrix.os }} - strategy: - matrix: - os: [macos-latest] - python-version: ['3.11', '3.12', '3.13'] - permissions: - contents: read - steps: - - name: Checkout Repository - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - name: Install Dependencies - run: | - brew update - brew install freetds - python -m pip install --upgrade pip - pip install pytest==9.0.2 lorem pytest-mock - pip install google-genai - pip install -r requirements.txt - pip install . - - name: Run Tests run: pytest env: @@ -92,46 +52,21 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - os: [ubuntu-latest, windows-latest, macos-14] - python-version: ['3.10', '3.11', '3.12', '3.13'] - permissions: - contents: read - steps: - - name: Checkout Repository - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Test pip install - run: | - python -m pip install --upgrade pip - pip install . - wrangles.recipe tests/samples/generate-data.wrgl.yml - - test-pip-install-macos-latest: - runs-on: ${{ matrix.os }} - strategy: - matrix: - os: [macos-latest] - python-version: ['3.11', '3.12', '3.13'] + os: [ubuntu-latest, windows-latest] + python-version: ['3.11', '3.13'] permissions: contents: read steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - name: Test pip install run: | - brew update - brew install freetds python -m pip install --upgrade pip pip install . wrangles.recipe tests/samples/generate-data.wrgl.yml @@ -140,42 +75,42 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: '3.12' - name: Install Dependencies run: | - pip install -r requirements.txt + pip install -r requirements-full.txt pip install jsonschema - name: Generate and Test Schema run: cd schema && python generate_recipe_schema.py - name: Save Schema as Artifact - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: schema path: schema/schema.json build: runs-on: ubuntu-latest - needs: [pytest, pytest-macos-latest, test-pip-install, test-pip-install-macos-latest, test-generate-schema] + needs: [pytest, test-pip-install, test-generate-schema] permissions: contents: read packages: write steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 # Login against a Docker registry except on PR # https://github.com/docker/login-action - name: Log into registry ${{ env.REGISTRY }} if: github.event_name != 'pull_request' - uses: docker/login-action@v3 + uses: docker/login-action@v4 with: registry: ${{ env.REGISTRY }} username: ${{ github.actor }} @@ -185,14 +120,14 @@ jobs: # https://github.com/docker/metadata-action - name: Extract Docker metadata id: meta - uses: docker/metadata-action@v5 + uses: docker/metadata-action@v6 with: images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} # Build and push Docker image with Buildx (don't push on PR) # https://github.com/docker/build-push-action - name: Build and push Docker image - uses: docker/build-push-action@v5 + uses: docker/build-push-action@v7 with: context: . push: ${{ github.event_name != 'pull_request' }} @@ -216,7 +151,7 @@ jobs: GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Remove wrangles folder run: rm -r wrangles diff --git a/.github/workflows/publish-tagged.yml b/.github/workflows/publish-tagged.yml index 3a3cf0884..9d4ce719d 100644 --- a/.github/workflows/publish-tagged.yml +++ b/.github/workflows/publish-tagged.yml @@ -16,18 +16,17 @@ jobs: contents: read steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - name: Install Dependencies run: | python -m pip install --upgrade pip pip install pytest==9.0.2 lorem pytest-mock - pip install google-genai pip install -r requirements.txt pip install . @@ -49,12 +48,12 @@ jobs: contents: read steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - name: Test pip install run: | @@ -62,14 +61,48 @@ jobs: pip install . wrangles.recipe tests/samples/generate-data.wrgl.yml - test-generate-schema: - runs-on: ubuntu-latest + pytest-macos: + runs-on: macos-latest + permissions: + contents: read steps: - name: Checkout Repository uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install Dependencies + run: | + brew update + brew install freetds + python -m pip install --upgrade pip + pip install pytest==9.0.2 lorem pytest-mock + pip install -r requirements.txt + pip install . + + - name: Run Tests + run: pytest + env: + WRANGLES_USER: ${{ secrets.WRANGLES_USER }} + WRANGLES_PASSWORD: ${{ secrets.WRANGLES_PASSWORD }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + HUGGINGFACE_TOKEN: ${{ secrets.HUGGINGFACE_TOKEN }} + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + SERPAPI_API_KEY: ${{ secrets.SERPAPI_API_KEY }} + GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} + + test-generate-schema: + runs-on: ubuntu-latest + steps: + - name: Checkout Repository + uses: actions/checkout@v5 + + - name: Set up Python + uses: actions/setup-python@v6 with: python-version: '3.12' @@ -82,20 +115,20 @@ jobs: run: cd schema && python generate_recipe_schema.py - name: Save Schema as Artifact - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: schema path: schema/schema.json build: runs-on: ubuntu-latest - needs: [pytest, test-pip-install, test-generate-schema] + needs: [pytest, test-pip-install, pytest-macos, test-generate-schema] permissions: contents: read packages: write steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Get release tag run: echo "RELEASE_VERSION=$(echo ${GITHUB_REF:11})" >> $GITHUB_ENV @@ -104,7 +137,7 @@ jobs: # https://github.com/docker/login-action - name: Log into registry ${{ env.REGISTRY }} if: github.event_name != 'pull_request' - uses: docker/login-action@v3 + uses: docker/login-action@v4 with: registry: ${{ env.REGISTRY }} username: ${{ github.actor }} @@ -114,14 +147,14 @@ jobs: # https://github.com/docker/metadata-action - name: Extract Docker metadata id: meta - uses: docker/metadata-action@v5 + uses: docker/metadata-action@v6 with: images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} # Build and push Docker image with Buildx (don't push on PR) # https://github.com/docker/build-push-action - name: Build and push Docker image - uses: docker/build-push-action@v5 + uses: docker/build-push-action@v7 with: context: . push: ${{ github.event_name != 'pull_request' }} @@ -136,12 +169,12 @@ jobs: # id-token: write # steps: # - name: Checkout Repository - # uses: actions/checkout@v4 + # uses: actions/checkout@v5 # - name: Set up Python - # uses: actions/setup-python@v5 + # uses: actions/setup-python@v6 # with: - # python-version: "3.10" + # python-version: "3.11" # - name: Install build # run: >- @@ -173,12 +206,12 @@ jobs: id-token: write steps: - name: Checkout Repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: - python-version: "3.10" + python-version: "3.11" - name: Install build run: >- diff --git a/.gitignore b/.gitignore index 897ca8f83..29aa9f44b 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,10 @@ tests/temp/* !tests/temp/README.md +# Local/generated output files +output_files/* +!output_files/.gitkeep + # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/README.md b/README.md index aea62bc03..1469b7d11 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,30 @@ The python package can be installed using [pip](https://pip.pypa.io/en/stable/ge pip install wrangles ``` +This installs the core package, which covers the vast majority of use cases: all data wrangles, recipe execution, Excel and CSV file I/O, HTTP connectors, and SQLite. + +### Optional dependencies + +Connectors for databases, cloud storage, and external services require additional packages. Install only the ones you need: + +| Capability | Install | +|---|---| +| Microsoft SQL Server | `pip install pymssql sqlalchemy` | +| Microsoft Access | `pip install pyodbc` | +| DuckDB | `pip install duckdb` | +| PostgreSQL | `pip install psycopg2-binary sqlalchemy` | +| MySQL | `pip install pymysql sqlalchemy` | +| MongoDB | `pip install pymongo[srv]` | +| AWS S3 | `pip install boto3` | +| Salesforce | `pip install simple-salesforce` | +| SFTP / SSH | `pip install fabric` | +| Notifications | `pip install apprise` | +| OpenAI SDK | `pip install openai` | +| Google Gemini | `pip install google-generativeai` | +| SerpAPI (web search) | `pip install serpapi` | + +> If a connector is used without its required package installed, Wrangles will raise a clear `ImportError` with the exact `pip install` command needed. + Once installed, import the package into your code. ```python import wrangles diff --git a/dockerfile b/dockerfile index 2bb573109..9bf6d006a 100644 --- a/dockerfile +++ b/dockerfile @@ -1,4 +1,4 @@ -FROM python:3.10.16-slim-bookworm AS compile-image +FROM python:3.11-slim-bookworm AS compile-image # Copy package COPY . /pkg @@ -21,15 +21,15 @@ RUN CFLAGS="-g0 -Wl,--strip-all" pip install --no-cache-dir --compile --global-o RUN pip install --no-cache-dir /pkg # Botocore contains lots of definitions for all AWS services. We are only using S3. Remove all other files to save space -RUN cd /opt/venv/lib/python3.10/site-packages/botocore/data && cp -r s3 _retry.json endpoints.json partitions.json sdk-default-configuration.json /tmp/ -RUN rm -r /opt/venv/lib/python3.10/site-packages/botocore/data/* -RUN cp -r /tmp/s3 /tmp/_retry.json /tmp/endpoints.json /tmp/partitions.json /tmp/sdk-default-configuration.json /opt/venv/lib/python3.10/site-packages/botocore/data +RUN cd /opt/venv/lib/python3.11/site-packages/botocore/data && cp -r s3 _retry.json endpoints.json partitions.json sdk-default-configuration.json /tmp/ +RUN rm -r /opt/venv/lib/python3.11/site-packages/botocore/data/* +RUN cp -r /tmp/s3 /tmp/_retry.json /tmp/endpoints.json /tmp/partitions.json /tmp/sdk-default-configuration.json /opt/venv/lib/python3.11/site-packages/botocore/data # Pandas contains a lot of unnecessary test data that we won't use -RUN rm -r /opt/venv/lib/python3.10/site-packages/pandas/tests/* +RUN rm -r /opt/venv/lib/python3.11/site-packages/pandas/tests/* # Create build image -FROM python:3.10.16-slim-bookworm AS build-image +FROM python:3.11-slim-bookworm AS build-image COPY --from=compile-image /opt/venv /opt/venv LABEL maintainer="WrangleWorks" diff --git a/output_files/.gitkeep b/output_files/.gitkeep new file mode 100644 index 000000000..e69de29bb diff --git a/schema/generate_recipe_schema.py b/schema/generate_recipe_schema.py index 49fe3da3e..01380ead6 100644 --- a/schema/generate_recipe_schema.py +++ b/schema/generate_recipe_schema.py @@ -158,24 +158,28 @@ def getMethodDocs(schema_wrangles, obj, path): # Add common wrangle properties for wrangle in schema['wrangles']: - if "properties" not in schema['wrangles'][wrangle]: - schema['wrangles'][wrangle]["properties"] = {} + if "properties" in schema['wrangles'][wrangle]: + wrangle_properties = schema['wrangles'][wrangle]['properties'] + else: + # Wrangles defined with a top-level anyOf (e.g. recipe, which is + # recursive) keep their properties in the final anyOf branch + wrangle_properties = schema['wrangles'][wrangle]['anyOf'][-1]['properties'] - schema['wrangles'][wrangle]['properties']["if"] = { + wrangle_properties["if"] = { "$ref": f"#/$defs/wrangles/commonProperties/if" } if wrangle not in wrangles.config.where_not_implemented: if wrangle in wrangles.config.where_overwrite_output: - schema['wrangles'][wrangle]['properties']['where'] = { + wrangle_properties['where'] = { "$ref": "#/$defs/wrangles/commonProperties/where_special" } else: - schema['wrangles'][wrangle]['properties']['where'] = { + wrangle_properties['where'] = { "$ref": "#/$defs/wrangles/commonProperties/where" } - schema['wrangles'][wrangle]['properties']["where_params"] = { + wrangle_properties["where_params"] = { "$ref": f"#/$defs/wrangles/commonProperties/where_params" } diff --git a/schema/recipe_base_schema.json b/schema/recipe_base_schema.json index 03f7dcca4..c3d196e53 100644 --- a/schema/recipe_base_schema.json +++ b/schema/recipe_base_schema.json @@ -63,6 +63,15 @@ "alias": { "type": "array", "description": "Placeholder to store YAML anchor values for use with aliases elsewhere in the recipe" + }, + "where": { + "$ref": "#/$defs/wrangles/commonProperties/where" + }, + "where_params": { + "$ref": "#/$defs/wrangles/commonProperties/where_params" + }, + "if": { + "$ref": "#/$defs/wrangles/commonProperties/if" } }, "$defs": { diff --git a/tests/connectors/test_access.py b/tests/connectors/test_access.py new file mode 100644 index 000000000..ec31891f4 --- /dev/null +++ b/tests/connectors/test_access.py @@ -0,0 +1,109 @@ +import pandas as pd +from types import SimpleNamespace +from unittest.mock import Mock + +from wrangles.connectors import access + + +class MockTables: + def __init__(self, exists): + self.exists = exists + + def fetchone(self): + return object() if self.exists else None + + +class MockCursor: + def __init__(self, table_exists=False): + self.table_exists = table_exists + self.executed = [] + self.executemany_calls = [] + + def tables(self, table=None, tableType=None): + return MockTables(self.table_exists) + + def execute(self, sql, params=()): + self.executed.append((sql, params)) + return self + + def executemany(self, sql, rows): + self.executemany_calls.append((sql, list(rows))) + return self + + +class MockConnection: + def __init__(self, cursor): + self.cursor_obj = cursor + self.committed = False + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def cursor(self): + return self.cursor_obj + + def commit(self): + self.committed = True + + +def test_connection_string_requires_database_or_connection_string(): + try: + access._connection_string() + assert False + except ValueError as e: + assert str(e) == 'database or connection_string must be provided' + + +def test_read_sql(monkeypatch): + data = pd.DataFrame({'Col1': ['Data1', 'Data2']}) + mock_connect = Mock(return_value=MockConnection(MockCursor())) + monkeypatch.setattr(access, "_pyodbc", SimpleNamespace(connect=mock_connect)) + monkeypatch.setattr(pd, "read_sql", Mock(return_value=data)) + + df = access.read( + database='database.accdb', + command='SELECT * from df_mock' + ) + + assert df.equals(data) + mock_connect.assert_called_once_with('DRIVER={Microsoft Access Driver (*.mdb, *.accdb)};DBQ=database.accdb;') + + +def test_write_sql_creates_table_and_inserts(monkeypatch): + cursor = MockCursor(table_exists=False) + monkeypatch.setattr(access, "_pyodbc", SimpleNamespace(connect=Mock(return_value=MockConnection(cursor)))) + + result = access.write( + df=pd.DataFrame({'Col1': ['Data1', 'Data2'], 'Col2': [1, 2]}), + database='database.accdb', + table='WrWx' + ) + + assert result is None + assert cursor.executed[0][0] == 'CREATE TABLE [WrWx] ([Col1] LONGTEXT, [Col2] INTEGER)' + assert cursor.executemany_calls[0] == ( + 'INSERT INTO [WrWx] ([Col1], [Col2]) VALUES (?, ?)', + [('Data1', 1), ('Data2', 2)] + ) + + +def test_run(monkeypatch): + cursor = MockCursor() + connection = MockConnection(cursor) + monkeypatch.setattr(access, "_pyodbc", SimpleNamespace(connect=Mock(return_value=connection))) + + result = access.run( + database='database.accdb', + command=['DELETE FROM WrWx WHERE Col1 = ?', 'UPDATE WrWx SET Col1 = ?'], + params=('Data1',) + ) + + assert result is None + assert cursor.executed == [ + ('DELETE FROM WrWx WHERE Col1 = ?', ('Data1',)), + ('UPDATE WrWx SET Col1 = ?', ('Data1',)) + ] + assert connection.committed diff --git a/tests/connectors/test_akeneo.py b/tests/connectors/test_akeneo.py new file mode 100644 index 000000000..3c222c8f7 --- /dev/null +++ b/tests/connectors/test_akeneo.py @@ -0,0 +1,104 @@ +import os +import pytest +import pandas as pd +from wrangles.connectors.akeneo import read, write + + +_host = os.getenv('AKENEO_HOST', '...') +_user = os.getenv('AKENEO_USERNAME', '...') +_password = os.getenv('AKENEO_PASSWORD', '...') +_client_id = os.getenv('AKENEO_CLIENT_ID', '...') +_client_secret = os.getenv('AKENEO_SECRET', '...') + +_has_credentials = all( + os.getenv(v) + for v in ('AKENEO_HOST', 'AKENEO_USERNAME', 'AKENEO_PASSWORD', 'AKENEO_CLIENT_ID', 'AKENEO_SECRET') +) + +_skip_no_creds = pytest.mark.skipif( + not _has_credentials, + reason="Akeneo credentials not set in environment variables" +) + + +@_skip_no_creds +def test_read(): + df = read( + host=_host, + user=_user, + password=_password, + client_id=_client_id, + client_secret=_client_secret, + source="products", + ) + assert isinstance(df, pd.DataFrame) + assert len(df) > 0 + + +@_skip_no_creds +def test_read_columns(): + df = read( + host=_host, + user=_user, + password=_password, + client_id=_client_id, + client_secret=_client_secret, + source="products", + columns=["identifier"], + ) + assert df.columns.tolist() == ["identifier"] + assert len(df) > 0 + + +@_skip_no_creds +def test_read_pagination(): + df = read( + host=_host, + user=_user, + password=_password, + client_id=_client_id, + client_secret=_client_secret, + source="products", + ) + assert isinstance(df, pd.DataFrame) + + +@_skip_no_creds +def test_read_auth_error(): + with pytest.raises(ValueError, match="Akeneo authentication failed"): + read( + host=_host, + user="bad_user", + password="bad_pass", + client_id="bad_client_id", + client_secret="bad_secret", + source="products", + ) + + +@_skip_no_creds +def test_read_api_error(): + with pytest.raises(ValueError, match="Status Code:"): + read( + host=_host, + user=_user, + password=_password, + client_id=_client_id, + client_secret=_client_secret, + source="invalid_source", + ) + + +@_skip_no_creds +def test_write_error(): + df = pd.DataFrame({"identifier": ["prod1"]}) + with pytest.raises(ValueError, match="Status Code:"): + write( + df=df, + host=_host, + user=_user, + password=_password, + client_id=_client_id, + client_secret=_client_secret, + source="invalid_source", + ) diff --git a/tests/connectors/test_concurrent.py b/tests/connectors/test_concurrent.py index 843e96a0b..8ca6dab8e 100644 --- a/tests/connectors/test_concurrent.py +++ b/tests/connectors/test_concurrent.py @@ -1,50 +1,108 @@ -import wrangles -from datetime import datetime import time + import pandas as pd +import wrangles + + +# Shared timing values for all concurrent timing tests. +# +# Each task waits for 5 seconds. If the tasks run concurrently, total runtime +# should be roughly 5 seconds plus overhead. If they run serially, total runtime +# should be roughly 15 seconds. +WAIT_SECONDS = 5 +TASK_COUNT = 3 +SERIAL_RUNTIME_SECONDS = WAIT_SECONDS * TASK_COUNT + +# Allow CI overhead, especially on Windows, while still proving the tasks did +# not simply run one after another. +MAX_CONCURRENT_SECONDS = SERIAL_RUNTIME_SECONDS - 2 # 13 seconds + + +def _assert_concurrent_runtime(elapsed): + """ + Assert that a set of WAIT_SECONDS tasks ran concurrently rather than serially. + """ + assert elapsed >= WAIT_SECONDS, ( + f"Expected runtime to be at least {WAIT_SECONDS}s because each task waits " + f"{WAIT_SECONDS}s, but elapsed time was {elapsed:.2f}s." + ) + + assert elapsed < MAX_CONCURRENT_SECONDS, ( + f"Expected concurrent runtime to be materially less than the " + f"{SERIAL_RUNTIME_SECONDS}s serial runtime, but elapsed time was " + f"{elapsed:.2f}s." + ) + + +def _run_recipe_and_measure(recipe, functions): + """ + Run a recipe and return both the recipe result and elapsed time. + + time.monotonic() is used for elapsed-time measurement because it is designed + to move only forward and is not affected by wall-clock changes. + """ + start = time.monotonic() + result = wrangles.recipe.run(recipe, functions=functions) + elapsed = time.monotonic() - start + + return result, elapsed + def wait_and_read(duration): time.sleep(duration) return pd.DataFrame({"column": [f"value{str(duration)}"]}) -def test_read_multithread(): + +def wait_and_append(df, wait, test_vals_key): + time.sleep(wait) + test_vals[test_vals_key].append(wait) + + +def sleep(seconds): + """ + As a builtin function implemented in C, time.sleep() is not directly + accessible to the recipe engine. Wrap it as a Python function to make the + function definition available. """ - Test using the concurrent connector to read multithreaded - """ - start = datetime.now() + time.sleep(seconds) - df = wrangles.recipe.run( - """ + +test_vals = { + "multithread": [], +} + + +def test_read_multithread(): + """ + Test using the concurrent connector to read using multiple threads. + """ + df, elapsed = _run_recipe_and_measure( + f""" read: - union: sources: - concurrent: read: - custom.wait_and_read: - duration: 5 + duration: {WAIT_SECONDS} - custom.wait_and_read: - duration: 2 + duration: {WAIT_SECONDS} - custom.wait_and_read: - duration: 3 + duration: {WAIT_SECONDS} """, - functions=wait_and_read + functions=wait_and_read, ) - end = datetime.now() + assert len(df) == TASK_COUNT + _assert_concurrent_runtime(elapsed) - assert ( - 5 <= (end - start).seconds < 10 and - len(df) == 3 - ) def test_read_multiprocess(): """ - Test using the concurrent connector to read using multiprocessing - """ - start = datetime.now() - - df = wrangles.recipe.run( - """ + Test using the concurrent connector to read using multiprocessing. + """ + df, elapsed = _run_recipe_and_measure( + f""" read: - union: sources: @@ -52,38 +110,27 @@ def test_read_multiprocess(): use_multiprocessing: true read: - custom.wait_and_read: - duration: 5 + duration: {WAIT_SECONDS} - custom.wait_and_read: - duration: 2 + duration: {WAIT_SECONDS} - custom.wait_and_read: - duration: 3 + duration: {WAIT_SECONDS} """, - functions=wait_and_read + functions=wait_and_read, ) - end = datetime.now() + assert len(df) == TASK_COUNT + _assert_concurrent_runtime(elapsed) - assert ( - 5 <= (end - start).seconds < 10 and - len(df) == 3 - ) - -test_vals = { - "multithread": [], - "multiprocess": [] -} -def wait_and_append(df, wait, test_vals_key): - time.sleep(wait) - test_vals[test_vals_key].append(wait) def test_write_multithread(): """ - Test using the concurrent connector to write + Test using the concurrent connector to write using multiple threads. """ - start = datetime.now() + test_vals["multithread"] = [] - wrangles.recipe.run( - """ + _, elapsed = _run_recipe_and_measure( + f""" read: - test: rows: 1 @@ -93,84 +140,64 @@ def test_write_multithread(): - concurrent: write: - custom.wait_and_append: - wait: 5 + wait: {WAIT_SECONDS} test_vals_key: multithread - custom.wait_and_append: - wait: 2 + wait: {WAIT_SECONDS} test_vals_key: multithread - custom.wait_and_append: - wait: 3 + wait: {WAIT_SECONDS} test_vals_key: multithread """, - functions=wait_and_append + functions=wait_and_append, ) - end = datetime.now() - - assert ( - 5 <= (end - start).seconds < 10 and - test_vals["multithread"] == [2,3,5] - ) + assert test_vals["multithread"] == [WAIT_SECONDS] * TASK_COUNT + _assert_concurrent_runtime(elapsed) -def sleep(seconds): - """ - As a builtin function implemented in C, - time.sleep() is not directly accessible to the recipe engine. - Wrap as a python function to make the function definition available. - """ - time.sleep(seconds) def test_run_multithread(): """ - Test using the concurrent connector to run - using multithreading + Test using the concurrent connector to run using multiple threads. """ - start = datetime.now() - - wrangles.recipe.run( - """ + _, elapsed = _run_recipe_and_measure( + f""" run: on_start: - concurrent: run: - custom.sleep: - seconds: 5 + seconds: {WAIT_SECONDS} - custom.sleep: - seconds: 2 + seconds: {WAIT_SECONDS} - custom.sleep: - seconds: 3 + seconds: {WAIT_SECONDS} """, - functions=sleep + functions=sleep, ) - end = datetime.now() + _assert_concurrent_runtime(elapsed) - assert 5 <= (end - start).seconds < 10 def test_run_multiprocess(): """ - Test using the concurrent connector to run - using multiprocessing + Test using the concurrent connector to run using multiprocessing. """ - start = datetime.now() - - wrangles.recipe.run( - """ + _, elapsed = _run_recipe_and_measure( + f""" run: on_start: - concurrent: use_multiprocessing: true run: - custom.sleep: - seconds: 5 + seconds: {WAIT_SECONDS} - custom.sleep: - seconds: 2 + seconds: {WAIT_SECONDS} - custom.sleep: - seconds: 3 + seconds: {WAIT_SECONDS} """, - functions=sleep + functions=sleep, ) - end = datetime.now() - - assert 5 <= (end - start).seconds < 10 \ No newline at end of file + _assert_concurrent_runtime(elapsed) \ No newline at end of file diff --git a/tests/connectors/test_duckdb.py b/tests/connectors/test_duckdb.py new file mode 100644 index 000000000..fa7272f4a --- /dev/null +++ b/tests/connectors/test_duckdb.py @@ -0,0 +1,72 @@ +import importlib.util + +import pandas as pd +import pytest + +from wrangles.connectors import duckdb + + +pytestmark = pytest.mark.skipif( + importlib.util.find_spec('duckdb') is None, + reason='duckdb optional dependency is not installed' +) + + +def test_read_sql(tmp_path): + database = tmp_path / 'test.duckdb' + duckdb.run( + database=str(database), + command='CREATE TABLE df_mock AS SELECT \'Data1\' AS Col1, 1 AS Col2 UNION ALL SELECT \'Data2\', 2' + ) + + df = duckdb.read( + database=str(database), + command='SELECT * from df_mock ORDER BY Col2' + ) + + assert df.equals( + pd.DataFrame( + { + 'Col1': ['Data1', 'Data2'], + 'Col2': pd.array([1, 2], dtype='int32') + } + ) + ) + + +def test_write_sql(tmp_path): + database = tmp_path / 'write.duckdb' + df = pd.DataFrame({'Col1': ['Data1', 'Data2']}) + + duckdb.write( + df=df, + database=str(database), + table='temp_mock' + ) + + assert duckdb.read( + database=str(database), + command='SELECT * from temp_mock' + ).equals(df) + + +def test_run(tmp_path): + database = tmp_path / 'run.duckdb' + df = pd.DataFrame({'Col1': ['Data1', 'Data2']}) + duckdb.write( + df=df, + database=str(database), + table='test_table' + ) + + duckdb.run( + database=str(database), + command='CREATE TABLE test_table_copy AS SELECT * FROM test_table' + ) + + df_copy = duckdb.read( + database=str(database), + command='SELECT * from test_table_copy' + ) + + assert df.equals(df_copy) diff --git a/tests/connectors/test_file.py b/tests/connectors/test_file.py index 7d9756e85..e187d4e4e 100644 --- a/tests/connectors/test_file.py +++ b/tests/connectors/test_file.py @@ -1,6 +1,7 @@ """ Test the file connector for reading and writing files to the local file system. """ +import pathlib import uuid as _uuid import wrangles import pandas as _pd @@ -953,6 +954,65 @@ def mock_write_excel(self, **kwargs): # col2: no override, so default top is applied assert captured['column_formats']['col2']['valign'] == 'top' + class TestIllegalCharacters: + """ + Tests for automatic removal of illegal XML characters when writing .xlsx files. + openpyxl raises IllegalCharacterError for chars outside the XML 1.0 legal range. + """ + + def test_write_xlsx_null_byte_no_error(self): + """ + Writing a DataFrame with null bytes should not raise IllegalCharacterError. + """ + filename = 'tests/temp/illegal_chars_null.xlsx' + df = _pd.DataFrame({'col': ['hello\x00world']}) + wrangles.connectors.file.write(df, name=filename) + result = _pd.read_excel(filename) + assert result['col'][0] == 'helloworld' + + def test_write_xlsx_control_chars_stripped(self): + """ + ASCII control characters (0x01–0x08, 0x0b, 0x0c, 0x0e–0x1f) are stripped. + """ + filename = 'tests/temp/illegal_chars_ctrl.xlsx' + df = _pd.DataFrame({'col': ['a\x01\x02\x03\x08\x0b\x0c\x0e\x1fb']}) + wrangles.connectors.file.write(df, name=filename) + result = _pd.read_excel(filename) + assert result['col'][0] == 'ab' + + def test_write_xlsx_allowed_whitespace_preserved(self): + """ + Tab (0x09), LF (0x0a), and CR (0x0d) are legal XML chars and must not be removed. + """ + filename = 'tests/temp/illegal_chars_whitespace.xlsx' + df = _pd.DataFrame({'col': ['a\tb\nc\rd']}) + wrangles.connectors.file.write(df, name=filename) + result = _pd.read_excel(filename) + assert result['col'][0] == 'a\tb\nc\rd' + + def test_write_xlsx_clean_data_unchanged(self): + """ + Normal strings must pass through unchanged (regression guard). + """ + filename = 'tests/temp/illegal_chars_clean.xlsx' + df = _pd.DataFrame({'col': ['hello', 'world', 'normal text']}) + wrangles.connectors.file.write(df, name=filename) + result = _pd.read_excel(filename) + assert result['col'].tolist() == ['hello', 'world', 'normal text'] + + def test_write_xlsx_in_memory_no_error(self): + """ + The in-memory (file_object) path is also covered. + """ + from io import BytesIO + buf = BytesIO() + df = _pd.DataFrame({'col': ['null\x00byte', 'ctrl\x01char']}) + wrangles.connectors.file.write(df, name='output.xlsx', file_object=buf) + buf.seek(0) + result = _pd.read_excel(buf) + assert result['col'].tolist() == ['nullbyte', 'ctrlchar'] + + def test_read_object(): """ Test reading a file passed directly into the recipe as an object @@ -993,3 +1053,42 @@ def test_read_object_json(): } ) assert len(df) == 3 and df['Col1'][0] == 'a' + + +class TestPathObjects: + """ + Test that read and write accept pathlib.Path objects (issue #986) + """ + def test_read_csv_path_object(self): + """ + Test that file.read accepts a pathlib.Path for a CSV file + """ + df = wrangles.connectors.file.read(pathlib.Path('tests/samples/data.csv')) + assert list(df.columns) == ['Find', 'Replace'] + + def test_read_excel_path_object(self): + """ + Test that file.read accepts a pathlib.Path for an Excel file + """ + df = wrangles.connectors.file.read(pathlib.Path('tests/samples/data.xlsx')) + assert list(df.columns) == ['Find', 'Replace'] + + def test_write_csv_path_object(self, tmp_path): + """ + Test that file.write accepts a pathlib.Path for a CSV file + """ + dest = tmp_path / 'out.csv' + df = _pd.DataFrame({'col1': ['a'], 'col2': ['b']}) + wrangles.connectors.file.write(df, dest) + result = wrangles.connectors.file.read(dest) + assert result['col1'][0] == 'a' + + def test_write_excel_path_object(self, tmp_path): + """ + Test that file.write accepts a pathlib.Path for an Excel file + """ + dest = tmp_path / 'out.xlsx' + df = _pd.DataFrame({'col1': ['a'], 'col2': ['b']}) + wrangles.connectors.file.write(df, dest) + result = wrangles.connectors.file.read(dest) + assert result['col1'][0] == 'a' diff --git a/tests/connectors/test_postgres.py b/tests/connectors/test_postgres.py index c742fd435..8bb38923e 100644 --- a/tests/connectors/test_postgres.py +++ b/tests/connectors/test_postgres.py @@ -103,13 +103,15 @@ def test_write_sql_experimental(mocker): def test_psql_insert_copy(mocker): - m = mocker.patch("wrangles.connectors.postgres.write") - m2 = mocker.patch("psycopg2.connect") - config = { - 'table': m, - 'conn': m2, - 'keys': 'keys_mock', - 'data_iter': 'data_iter_mock' - } - df = _psql_insert_copy(**config) - assert df == None \ No newline at end of file + conn = mocker.MagicMock() + table = mocker.MagicMock() + table.schema = None + table.name = "WrWx" + result = _psql_insert_copy( + table=table, + conn=conn, + keys=["col1", "col2"], + data_iter=[("a", "b")], + ) + assert result is None + conn.connection.cursor().__enter__().copy_expert.assert_called_once() \ No newline at end of file diff --git a/tests/connectors/test_s3.py b/tests/connectors/test_s3.py index c851211fd..8d367ad3f 100644 --- a/tests/connectors/test_s3.py +++ b/tests/connectors/test_s3.py @@ -1,5 +1,6 @@ import pandas as pd import os +import pathlib import wrangles import pytest import time @@ -283,10 +284,10 @@ def test_upload_error_save_as_file_key(self): on_start: - s3.upload_files: bucket: wrwx-public - save_as: + file: - tests/samples/data.csv - tests/samples/data.json - file_key: + save_as: - Test_Upload_File.csv """ ) @@ -317,8 +318,56 @@ def test_run_upload_error_invalid_bucket_file_key(self): on_start: - s3.upload_files: bucket: wrwx-does-not-exist - file_key: does_not_exist.csv - save_as: tests/samples/data.csv + save_as: does_not_exist.csv + file: tests/samples/data.csv + """ + ) + + def test_run_upload_error_invalid_bucket_file_only(self): + """ + Gen 4 (and Gen 1): file only (no key/save_as) with an invalid bucket — + error comes from S3, confirming the param was accepted and auto-key + derivation ran before the upload attempt. + """ + with pytest.raises(RuntimeError, match="Failed to write"): + wrangles.recipe.run( + """ + run: + on_start: + - s3.upload_files: + bucket: wrwx-does-not-exist + file: tests/samples/data.csv + """ + ) + + def test_upload_error_missing_file(self): + """ + Gen 4: omitting file raises ValueError before reaching S3. + """ + with pytest.raises(ValueError, match="file must be provided"): + wrangles.recipe.run( + """ + run: + on_start: + - s3.upload_files: + bucket: wrwx-public + save_as: s3/key.csv + """ + ) + + def test_upload_error_missing_file_key_only(self): + """ + Gen 1 deprecated key kwarg: providing key alone (no file) raises + ValueError before reaching S3. + """ + with pytest.raises(ValueError, match="file must be provided"): + wrangles.recipe.run( + """ + run: + on_start: + - s3.upload_files: + bucket: wrwx-public + key: s3/key.csv """ ) @@ -393,8 +442,8 @@ def test_file_upload_and_download_file_key(self): on_start: - s3.upload_files: bucket: wrwx-public - file_key: Test_Upload_File.csv - save_as: tests/samples/data.csv + save_as: Test_Upload_File.csv + file: tests/samples/data.csv aws_access_key_id: {s3_key} aws_secret_access_key: {s3_secret} """ @@ -418,6 +467,62 @@ def test_file_upload_and_download_file_key(self): df = wrangles.recipe.run(recipe2) assert df.iloc[0]['Find'] == 'BRG' + def test_upload_pathlib_path_file(self): + """ + file param accepts a pathlib.Path object — RuntimeError from invalid + bucket confirms the Path was normalised and upload was attempted. + """ + with pytest.raises(RuntimeError, match="Failed to write"): + wrangles.connectors.s3.upload_files.run( + bucket='wrwx-does-not-exist', + file=pathlib.Path('tests/samples/data.csv'), + ) + + def test_upload_pathlib_path_file_list(self): + """ + file param accepts a list of pathlib.Path objects. + """ + with pytest.raises(RuntimeError, match="Failed to write"): + wrangles.connectors.s3.upload_files.run( + bucket='wrwx-does-not-exist', + file=[pathlib.Path('tests/samples/data.csv')], + ) + + def test_upload_bc_gen1_key_kwarg_upload_and_download(self): + """ + Gen 1 deprecated key kwarg: full upload + download cycle confirms + key is routed to save_as and the correct S3 object is written. + """ + wrangles.recipe.run( + f""" + run: + on_start: + - s3.upload_files: + bucket: wrwx-public + file: tests/samples/data.csv + key: Test_BC_gen1_key.csv + aws_access_key_id: {s3_key} + aws_secret_access_key: {s3_secret} + """ + ) + time.sleep(3) + df = wrangles.recipe.run( + f""" + run: + on_start: + - s3.download_files: + bucket: wrwx-public + file_key: Test_BC_gen1_key.csv + save_as: tests/temp/test_bc_gen1_key.csv + aws_access_key_id: {s3_key} + aws_secret_access_key: {s3_secret} + read: + - file: + name: tests/temp/test_bc_gen1_key.csv + """ + ) + assert df.iloc[0]['Find'] == 'BRG' + class TestRunDownload: """ Test using run s3.download_files @@ -518,4 +623,4 @@ def test_download_error_file_key(self): save_as: - tests/temp/temp_download_data.csv """ - ) \ No newline at end of file + ) diff --git a/tests/connectors/test_train.py b/tests/connectors/test_train.py index a05a86aff..1367135e4 100644 --- a/tests/connectors/test_train.py +++ b/tests/connectors/test_train.py @@ -1,7 +1,5 @@ import uuid -from pytest_mock import mocker - import wrangles import pandas as pd import pytest @@ -1361,32 +1359,33 @@ def test_upsert_key_only(self): assert 'Blade Runner Upsert' in df['City'].values assert 'New Value' in df['City'].values - def test_missing_columns_error_message(self, mocker): - """ - Verify that INSERT/UPSERT/UPDATE raise the expected error - when incoming columns are not present in the existing model. - """ + def test_missing_columns_error_message(self, mocker): + """ + Verify that INSERT/UPSERT/UPDATE raise the expected error + when incoming columns are not present in the existing model. + """ + + df = pd.DataFrame({ + "Key": ["k3"], + "Value": ["v3"], + "ExtraCol": ["x"] # This column does not exist in the model + }) + + # Mock the existing model so the test does not depend on a real, live model mocker.patch( - "wrangles.connectors.train._data.model_content", + "wrangles.data.model_content", return_value={ "Columns": ["Key", "Value"], - "Data": [["k1", "v1"]], - }, + "Data": [["k1", "v1"]] + } ) mocker.patch( - "wrangles.connectors.train._data.model", - return_value={"variant": "key"}, + "wrangles.data.model", + return_value={"variant": "key"} ) - df = pd.DataFrame({ - "Key": ["k3"], - "Value": ["v3"], - "ExtraCol": ["x"] # This column does not exist in the model - }) - - - # Test each action that performs the column-alignment check - for action in ("insert", "upsert", "update"): + # Test each action that performs the column-alignment check + for action in ("insert", "upsert", "update"): recipe = f""" write: - train.lookup: @@ -1682,3 +1681,120 @@ def test_meta_data_write_invalid_type_settings(self): """, dataframe=pd.DataFrame([{"settings": "not-a-dict"}]) ) + + +# +# Delete +# +class TestTrainDelete: + """ + Tests for wrangles.train.delete + """ + + def _get_model_id(self, response): + body = response.json() + model_id = body.get('model_id') or body.get('id') or body.get('modelId') or body.get('model') + assert model_id, f"No model_id in creation response: {body}" + return model_id + + def _mock_delete_ok(self, mocker): + mock = mocker.patch('wrangles.train._requests.delete') + mock.return_value.ok = True + mock.return_value.status_code = 200 + return mock + + def test_create_and_delete_classify(self, mocker): + """ + Train a new classify model (real API) then delete it (mocked). + Verifies creation succeeds and delete is called with the correct model_id. + """ + training_data = [ + ['apple', 'fruit', ''], + ['banana', 'fruit', ''], + ['carrot', 'vegetable', ''], + ] + response = wrangles.train.classify(training_data, name="Test Delete Classify Model") + assert response.ok, f"Model creation failed: {response.status_code} {response.text}" + + model_id = self._get_model_id(response) + mock_delete = self._mock_delete_ok(mocker) + + wrangles.train.delete(model_id) + + mock_delete.assert_called_once() + assert model_id in str(mock_delete.call_args) + + def test_create_and_delete_extract(self, mocker): + """ + Train a new extract model (real API) then delete it (mocked). + Verifies creation succeeds and delete is called with the correct model_id. + """ + training_data = [ + ['Television', 'TV', ''], + ['Refrigerator', 'Fridge', ''], + ['Automobile', 'Car', ''], + ] + response = wrangles.train.extract(training_data, name="Test Delete Extract Model") + assert response.ok, f"Model creation failed: {response.status_code} {response.text}" + + model_id = self._get_model_id(response) + mock_delete = self._mock_delete_ok(mocker) + + wrangles.train.delete(model_id) + + mock_delete.assert_called_once() + assert model_id in str(mock_delete.call_args) + + def test_create_and_delete_lookup(self, mocker): + """ + Train a new lookup model (real API) then delete it (mocked). + Verifies creation succeeds and delete is called with the correct model_id. + """ + data = [ + ['Key', 'Value'], + ['apple', 'fruit'], + ['carrot', 'vegetable'], + ] + response = wrangles.train.lookup(data, name="Test Delete Lookup Model", settings={'variant': 'key'}) + assert response.ok, f"Model creation failed: {response.status_code} {response.text}" + + model_id = self._get_model_id(response) + mock_delete = self._mock_delete_ok(mocker) + + wrangles.train.delete(model_id) + + mock_delete.assert_called_once() + assert model_id in str(mock_delete.call_args) + + def test_create_and_delete_standardize(self, mocker): + """ + Train a new standardize model (real API) then delete it (mocked). + Verifies creation succeeds and delete is called with the correct model_id. + """ + training_data = [ + ['ASAP', 'As Soon As Possible', ''], + ['ETA', 'Estimated Time of Arrival', ''], + ['USA', 'United States of America', ''], + ] + response = wrangles.train.standardize(training_data, name="Test Delete Standardize Model") + assert response.ok, f"Model creation failed: {response.status_code} {response.text}" + + model_id = self._get_model_id(response) + mock_delete = self._mock_delete_ok(mocker) + + wrangles.train.delete(model_id) + + mock_delete.assert_called_once() + assert model_id in str(mock_delete.call_args) + + def test_delete_error_raises_runtime_error(self, mocker): + """ + A non-OK response from the delete endpoint raises RuntimeError. + """ + mock = mocker.patch('wrangles.train._requests.delete') + mock.return_value.ok = False + mock.return_value.status_code = 404 + mock.return_value.text = '{"message":"Not Found"}' + + with pytest.raises(RuntimeError, match="Delete model failed"): + wrangles.train.delete("00000000-0000-0000") \ No newline at end of file diff --git a/tests/recipes/test_run.py b/tests/recipes/test_run.py index 24449ff5f..c50835246 100644 --- a/tests/recipes/test_run.py +++ b/tests/recipes/test_run.py @@ -5,6 +5,7 @@ in a test file for the respective connectors e.g. tests/connectors/test_notifications.py """ +import pathlib import pandas as pd import wrangles import pytest @@ -166,3 +167,19 @@ def run(key): ) assert check_var.get("value") is True + + +def test_run_path_object(): + """ + Test that recipe.run accepts a pathlib.Path to a YAML file (issue #986) + """ + df = wrangles.recipe.run(pathlib.Path('tests/samples/recipe-basic.wrgl.yml')) + assert list(df.columns) == ['header1', 'header2'] + + +def test_run_pure_posix_path(): + """ + Test that recipe.run accepts a pathlib.PurePosixPath (issue #986) + """ + df = wrangles.recipe.run(pathlib.PurePosixPath('tests/samples/recipe-basic.wrgl.yml')) + assert list(df.columns) == ['header1', 'header2'] diff --git a/tests/recipes/test_wrangles.py b/tests/recipes/test_wrangles.py index 4bf020705..f4e03a4cf 100644 --- a/tests/recipes/test_wrangles.py +++ b/tests/recipes/test_wrangles.py @@ -560,11 +560,64 @@ def test_if_condition_logging(self, caplog): # Check that only the executed wrangle appears in logs log_messages = [msg for msg in caplog.messages if "Wrangling ::" in msg] - assert any(": Wrangling :: convert.case :: col1 >> should_be_logged" in msg for msg in log_messages) - assert not any(": Wrangling :: convert.case :: col1 >> should_not_be_logged" in msg for msg in log_messages) + assert any(": Wrangling :: convert.case :: Completed :: col1 >> should_be_logged" in msg for msg in log_messages) + assert not any(": Wrangling :: convert.case :: Completed :: col1 >> should_not_be_logged" in msg for msg in log_messages) assert df['should_be_logged'][0] == 'VALUE' assert 'should_not_be_logged' not in df.columns + def test_if_false_with_missing_input_column(self): + """ + Test that a wrangle with a false if condition does not + raise a KeyError when its input column doesn't exist. + Column validation must not run before the if check. + """ + df = wrangles.recipe.run( + """ + read: + - test: + rows: 1 + values: + header: value + wrangles: + - merge.coalesce: + if: '"New_Column" in columns' + input: + - New_Column + - header + output: My Output + """ + ) + assert "My Output" not in df.columns + + def test_if_columns_check_then_conditional_wrangle(self): + """ + Test the full pattern from issue #765: first wrangle creates a + column conditionally, second wrangle uses that column only if it + was created — both guarded by if conditions checking 'columns'. + """ + df = wrangles.recipe.run( + """ + read: + - test: + rows: 1 + values: + header: value + wrangles: + - create.column: + if: '"NOT THERE" in columns' + output: New_Column + value: This column is new + + - merge.coalesce: + if: '"New_Column" in columns' + input: + - New_Column + - header + output: My Output + """ + ) + assert "My Output" not in df.columns + class TestPositionInput: """ Test using column indexes rather than names for input diff --git a/tests/recipes/wrangles/test_convert.py b/tests/recipes/wrangles/test_convert.py index 89d24d9e6..f1fe09f9c 100644 --- a/tests/recipes/wrangles/test_convert.py +++ b/tests/recipes/wrangles/test_convert.py @@ -252,6 +252,235 @@ def test_empty(self): ) assert df.empty and df.columns.to_list() == ['column', 'upper column'] + def test_sentence_tabs_after_punctuation(self): + """ + Sentence case with tab characters as whitespace between sentences. + The regex handles [ \t]* between punctuation and next word. + """ + data = pd.DataFrame({'Data': ["hello.\tthere! one more sentence."]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "Hello.\tThere! One more sentence." + + def test_sentence_leading_whitespace(self): + """ + Sentence case where the string starts with leading spaces/tabs. + The first non-whitespace character should be capitalised. + """ + data = pd.DataFrame({'Data': [" hello world. next sentence", "\t\thello again"]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == " Hello world. Next sentence" + assert df.iloc[1]['out'] == "\t\tHello again" + + def test_sentence_unicode_accented_chars(self): + """ + Sentence case preserves/capitalises Unicode accented characters. + """ + data = pd.DataFrame({'Data': ["hĆ©llo wƶrld. ƱoƱo is here! über cool."]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "HĆ©llo wƶrld. ƑoƱo is here! Über cool." + + def test_sentence_starts_with_number(self): + """ + Sentence case when the sentence starts with a digit. + The digit 'consumes' the capitalise flag; subsequent lowercase letters stay lower. + """ + data = pd.DataFrame({'Data': ["13 items found. that's all."]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "13 items found. That's all." + + def test_lower_non_string_preserved(self): + """ + Non-string values (int, list, None) in a lower-case column are returned unchanged. + The vectorised path must restore originals for non-string rows. + """ + data = pd.DataFrame({'Data': ["HELLO", 99, ["A", "B"], None]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: lower + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "hello" + assert df.iloc[1]['out'] == 99 + assert df.iloc[2]['out'] == ["A", "B"] + + def test_title_non_string_preserved(self): + """ + Non-string values in a title-case column are returned unchanged. + """ + data = pd.DataFrame({'Data': ["hello world", {"key": "val"}]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: title + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "Hello World" + assert df.iloc[1]['out'] == {"key": "val"} + + def test_purely_numeric_column_preserved(self): + """ + A column with an all-integer int64 dtype must not raise AttributeError + from the .str accessor and must return all values unchanged with a warning. + """ + data = pd.DataFrame({'Data': [1, 2, 3]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: lower + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == 1 + assert df.iloc[1]['out'] == 2 + assert df.iloc[2]['out'] == 3 + + def test_lower_large_dataframe(self): + """ + Vectorised lower case on a large dataframe produces correct output. + """ + n = 100_000 + data = pd.DataFrame({'Data': ['Hello World MIXED Case'] * n}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: lower + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Data'] == 'hello world mixed case' + assert df.iloc[-1]['Data'] == 'hello world mixed case' + + def test_upper_large_dataframe(self): + """ + Vectorised upper case on a large dataframe produces correct output. + """ + n = 100_000 + data = pd.DataFrame({'Data': ['hello world mixed case'] * n}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: upper + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Data'] == 'HELLO WORLD MIXED CASE' + + def test_title_large_dataframe(self): + """ + Vectorised title case on a large dataframe produces correct output. + """ + n = 100_000 + data = pd.DataFrame({'Data': ['hello world mixed case'] * n}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: title + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Data'] == 'Hello World Mixed Case' + + def test_sentence_large_dataframe(self): + """ + Regex-based sentence case on a large dataframe produces correct output. + """ + n = 50_000 + data = pd.DataFrame({'Data': ['first sentence. second one! third? yes.'] * n}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Data'] == 'First sentence. Second one! Third? Yes.' + + def test_lower_multiple_columns(self): + """ + Vectorised lower case across multiple input/output column pairs. + """ + n = 50_000 + data = pd.DataFrame({ + 'Col1': ['Hello World'] * n, + 'Col2': ['ANOTHER STRING'] * n, + 'Col3': ['YET ANOTHER'] * n, + }) + recipe = """ + wrangles: + - convert.case: + input: + - Col1 + - Col2 + - Col3 + output: + - Out1 + - Out2 + - Out3 + case: lower + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Out1'] == 'hello world' + assert df.iloc[0]['Out2'] == 'another string' + assert df.iloc[0]['Out3'] == 'yet another' + + def test_mixed_types_warning_logged_once(self, caplog): + """ + When a column contains non-string values, the invalid_data warning fires + exactly once regardless of how many non-string rows there are. + """ + import logging + data = pd.DataFrame({'Data': ["hello", 1, 2, 3, "world"]}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: upper + """ + with caplog.at_level(logging.WARNING): + df = wrangles.recipe.run(recipe, dataframe=data) + warning_messages = [r.message for r in caplog.records if 'invalid' in r.message.lower() or 'non-string' in r.message.lower() or 'not a string' in r.message.lower()] + # Warning should fire at most once + assert len(warning_messages) <= 1 + # Strings are correctly transformed + assert df.iloc[0]['Data'] == 'HELLO' + assert df.iloc[4]['Data'] == 'WORLD' + # Non-strings are preserved + assert df.iloc[1]['Data'] == 1 + class TestConvertDataType: """ @@ -884,6 +1113,142 @@ def test_empty(self): ) assert df.empty and df.columns.to_list() == ['header1', 'output column'] + def test_default_list_single_column(self): + """ + Test that a list of one default applies correctly when input is a list of one column + """ + df = wrangles.recipe.run( + """ + wrangles: + - convert.from_json: + input: + - header1 + default: + - {} + output: + - out1 + """, + dataframe=pd.DataFrame({ + "header1": ["", '[1,2,3]'] + }) + ) + assert df["out1"][0] == {} and df["out1"][1] == [1, 2, 3] + + def test_default_list_multiple_columns(self): + """ + Test that a list of defaults applies one default per column + """ + df = wrangles.recipe.run( + """ + wrangles: + - convert.from_json: + input: + - header1 + - header2 + default: + - {} + - [] + output: + - out1 + - out2 + """, + dataframe=pd.DataFrame({ + "header1": ["", '{"a": 1}'], + "header2": ["", '[4,5,6]'] + }) + ) + assert ( + df["out1"][0] == {} and + df["out1"][1] == {"a": 1} and + df["out2"][0] == [] and + df["out2"][1] == [4, 5, 6] + ) + + def test_default_list_one_applied_to_all_columns(self): + """ + Test that a list of one default is unwrapped and applied to all columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - convert.from_json: + input: + - header1 + - header2 + default: + - {} + output: + - out1 + - out2 + """, + dataframe=pd.DataFrame({ + "header1": ["", '{"a": 1}'], + "header2": ["", '{"b": 2}'] + }) + ) + assert ( + df["out1"][0] == {} and + df["out1"][1] == {"a": 1} and + df["out2"][0] == {} and + df["out2"][1] == {"b": 2} + ) + + def test_default_list_length_mismatch_raises(self): + """ + Test that a default list whose length doesn't match input raises an error + """ + with pytest.raises(ValueError, match="same length"): + wrangles.recipe.run( + """ + wrangles: + - convert.from_json: + input: + - header1 + - header2 + default: + - {} + - [] + - null + output: + - out1 + - out2 + """, + dataframe=pd.DataFrame({ + "header1": [""], + "header2": [""] + }) + ) + + def test_default_list_complex_values_no_output(self): + """ + Test a list of complex defaults (dict and list) with no output specified, + overwriting the input columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - convert.from_json: + input: + - header1 + - header2 + default: + - my_output: + key1: Val1 + key2: Val2 + - [this, is, a, list] + """, + dataframe=pd.DataFrame({ + "header1": ["", '{"a": 1}'], + "header2": ["", '[1,2,3]'] + }) + ) + assert ( + df["header1"][0] == {"my_output": {"key1": "Val1", "key2": "Val2"}} and + df["header1"][1] == {"a": 1} and + df["header2"][0] == ["this", "is", "a", "list"] and + df["header2"][1] == [1, 2, 3] + ) + class TestConvertFromYAML: """ @@ -1017,6 +1382,225 @@ def test_empty(self): ) assert df.empty and df.columns.to_list() == ['column', 'output column'] + def test_default_list_single_column(self): + """ + Test that a list of one default applies correctly when input is a list of one column + """ + df = wrangles.recipe.run( + """ + wrangles: + - convert.from_yaml: + input: + - column + default: + - {} + output: + - out1 + """, + dataframe=pd.DataFrame({ + "column": ["", "key: val\n"] + }) + ) + assert df["out1"][0] == {} and df["out1"][1] == {"key": "val"} + + def test_default_list_multiple_columns(self): + """ + Test that a list of defaults applies one default per column + """ + df = wrangles.recipe.run( + """ + wrangles: + - convert.from_yaml: + input: + - col1 + - col2 + default: + - {} + - [] + output: + - out1 + - out2 + """, + dataframe=pd.DataFrame({ + "col1": ["", "key: val\n"], + "col2": ["", "- item1\n"] + }) + ) + assert ( + df["out1"][0] == {} and + df["out1"][1] == {"key": "val"} and + df["out2"][0] == [] and + df["out2"][1] == ["item1"] + ) + + def test_default_list_one_applied_to_all_columns(self): + """ + Test that a list of one default is unwrapped and applied to all columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - convert.from_yaml: + input: + - col1 + - col2 + default: + - {} + output: + - out1 + - out2 + """, + dataframe=pd.DataFrame({ + "col1": ["", "key: val\n"], + "col2": ["", "key2: val2\n"] + }) + ) + assert ( + df["out1"][0] == {} and + df["out1"][1] == {"key": "val"} and + df["out2"][0] == {} and + df["out2"][1] == {"key2": "val2"} + ) + + def test_default_list_length_mismatch_raises(self): + """ + Test that a default list whose length doesn't match input raises an error + """ + with pytest.raises(ValueError, match="same length"): + wrangles.recipe.run( + """ + wrangles: + - convert.from_yaml: + input: + - col1 + - col2 + default: + - {} + - [] + - null + output: + - out1 + - out2 + """, + dataframe=pd.DataFrame({ + "col1": [""], + "col2": [""] + }) + ) + + def test_default_list_complex_values_no_output(self): + """ + Test a list of complex defaults (dict and list) with no output specified, + overwriting the input columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - convert.from_yaml: + input: + - col1 + - col2 + default: + - my_output: + key1: Val1 + key2: Val2 + - [this, is, a, list] + """, + dataframe=pd.DataFrame({ + "col1": ["", "a: 1\n"], + "col2": ["", "- 1\n- 2\n- 3\n"] + }) + ) + assert ( + df["col1"][0] == {"my_output": {"key1": "Val1", "key2": "Val2"}} and + df["col1"][1] == {"a": 1} and + df["col2"][0] == ["this", "is", "a", "list"] and + df["col2"][1] == [1, 2, 3] + ) + + def test_from_yaml_to_yaml_roundtrip_reported_data(self): + """ + Regression test for a reported round-trip bug (PR #987 review): + convert.from_yaml -> convert.to_yaml on multiple wildcard-matched + columns with per-column defaults should be idempotent - running + from_yaml again on the round-tripped output must reproduce the + same parsed values, for every row, including a row with a + malformed/empty cell that must fall back to its column's default. + """ + df = pd.DataFrame({ + "Top 1": [ + "Score: 0.295\nValue: Blade Runner", + "Score: 0.234\nValue: Interstellar", + "Score: 0.291\nValue: Interstellar", + ], + "Top 2": [ + "[]", + "Score: 0.22\nValue: Westworld", + "Score: 0.248\nValue: Blade Runner", + ], + "Top 3": [ + "Score: 0.174\nValue: Westworld", + "'", + "Score: 0.195\nValue: Westworld", + ], + }) + recipe = """ + wrangles: + - convert.from_yaml: + input: + - Top * + default: + - {} + - [] + - '' + """ + parsed = wrangles.recipe.run(recipe, dataframe=df.copy()) + + assert parsed.to_dict(orient="records") == [ + { + "Top 1": {"Score": 0.295, "Value": "Blade Runner"}, + "Top 2": [], + "Top 3": {"Score": 0.174, "Value": "Westworld"}, + }, + { + "Top 1": {"Score": 0.234, "Value": "Interstellar"}, + "Top 2": {"Score": 0.22, "Value": "Westworld"}, + # Malformed cell ("'") fails to parse and falls back to the default + "Top 3": "", + }, + { + "Top 1": {"Score": 0.291, "Value": "Interstellar"}, + "Top 2": {"Score": 0.248, "Value": "Blade Runner"}, + "Top 3": {"Score": 0.195, "Value": "Westworld"}, + }, + ] + + roundtripped = wrangles.recipe.run( + """ + wrangles: + - convert.from_yaml: + input: + - Top * + default: + - {} + - [] + - '' + - convert.to_yaml: + input: + - Top * + - convert.from_yaml: + input: + - Top * + default: + - {} + - [] + - '' + """, + dataframe=df.copy() + ) + + assert parsed.to_dict(orient="records") == roundtripped.to_dict(orient="records") + class TestConvertToJSON: """ diff --git a/tests/recipes/wrangles/test_create.py b/tests/recipes/wrangles/test_create.py index b7404c2ec..684871847 100644 --- a/tests/recipes/wrangles/test_create.py +++ b/tests/recipes/wrangles/test_create.py @@ -92,44 +92,165 @@ def test_create_columns_5(self): df = wrangles.recipe.run(recipe, dataframe=data) assert df.iloc[0]['column3'] in [True, False] - def test_column_exists(self): + def test_column_exists_no_error(self): """ - Check error if trying to create a column that already exists + Default behaviour: creating a column that already exists should not raise + an error — the existing column is left unchanged. """ - data = pd.DataFrame({ - 'col': ['data1'] - }) + data = pd.DataFrame({'col': ['data1']}) recipe = """ wrangles: - create.column: output: col + value: new_value """ - with pytest.raises(ValueError) as info: - wrangles.recipe.run(recipe, dataframe=data) - assert ( - info.typename == 'ValueError' and - '"col" column already exists in dataFrame.' in info.value.args[0] - ) + df = wrangles.recipe.run(recipe, dataframe=data) + assert df['col'][0] == 'data1' - def test_column_exists_list(self): + def test_column_exists_list_no_error(self): """ - Check error if trying to create a list of columns where one already exists + Default behaviour: creating a list of columns where one already exists + should not raise an error — the existing column is left unchanged. """ - data = pd.DataFrame({ - 'col': ['data1'] - }) + data = pd.DataFrame({'col': ['data1']}) recipe = """ wrangles: - create.column: output: - col + - col2 + value: new_value """ - with pytest.raises(ValueError) as info: + df = wrangles.recipe.run(recipe, dataframe=data) + assert df['col'][0] == 'data1' and df['col2'][0] == 'new_value' + + def test_column_exists_value_if_exists_existing(self): + """ + Explicit value_if_exists: existing leaves the column unchanged. + """ + data = pd.DataFrame({'col': ['data1']}) + recipe = """ + wrangles: + - create.column: + output: col + value: new_value + value_if_exists: existing + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df['col'][0] == 'data1' + + def test_column_exists_value_if_exists_new(self): + """ + value_if_exists: new overwrites the entire existing column. + """ + data = pd.DataFrame({'col': ['data1', 'data2']}) + recipe = """ + wrangles: + - create.column: + output: col + value: replaced + value_if_exists: new + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == ['replaced', 'replaced'] + + def test_column_exists_value_if_exists_coalesce(self): + """ + value_if_exists: coalesce fills only null/empty cells, leaving non-null + cells intact. + """ + data = pd.DataFrame({'col': ['keep', None, '']}) + recipe = """ + wrangles: + - create.column: + output: col + value: filled + value_if_exists: coalesce + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == ['keep', 'filled', 'filled'] + + def test_column_exists_value_if_exists_coalesce_numeric(self): + """ + value_if_exists: coalesce on a numeric column only fills NaN cells, + without comparing values to an empty string. + """ + data = pd.DataFrame({'col': [1, None, 3]}) + recipe = """ + wrangles: + - create.column: + output: col + value: 99 + value_if_exists: coalesce + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == [1, 99, 3] + + def test_column_exists_coalesce_value_new(self): + """ + coalesce_value: new prefers the new value over a non-empty existing + value, only falling back to the existing value where the new value + is empty/null. + """ + data = pd.DataFrame({'col': ['keep1', 'keep2']}) + recipe = """ + wrangles: + - create.column: + output: col + value: new_value + value_if_exists: coalesce + coalesce_value: new + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == ['new_value', 'new_value'] + + def test_column_exists_coalesce_value_new_fallback(self): + """ + coalesce_value: new falls back to the existing value when the new + value is empty/null. + """ + data = pd.DataFrame({'col': ['keep1', 'keep2']}) + recipe = """ + wrangles: + - create.column: + output: col + value: + value_if_exists: coalesce + coalesce_value: new + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == ['keep1', 'keep2'] + + def test_column_exists_coalesce_value_invalid(self): + """ + An invalid coalesce_value option raises a ValueError. + """ + data = pd.DataFrame({'col': ['data1']}) + recipe = """ + wrangles: + - create.column: + output: col + value: new_value + value_if_exists: coalesce + coalesce_value: not_a_real_option + """ + with pytest.raises(ValueError, match="coalesce_value"): + wrangles.recipe.run(recipe, dataframe=data) + + def test_column_exists_value_if_exists_invalid(self): + """ + An invalid value_if_exists option raises a ValueError. + """ + data = pd.DataFrame({'col': ['data1']}) + recipe = """ + wrangles: + - create.column: + output: col + value: new_value + value_if_exists: existng + """ + with pytest.raises(ValueError, match="value_if_exists"): wrangles.recipe.run(recipe, dataframe=data) - assert ( - info.typename == 'ValueError' and - "['col'] column(s)" in info.value.args[0] - ) def test_create_column_value_number(self): """ diff --git a/tests/recipes/wrangles/test_extract.py b/tests/recipes/wrangles/test_extract.py index 5e0e7ae74..87e6e7a74 100644 --- a/tests/recipes/wrangles/test_extract.py +++ b/tests/recipes/wrangles/test_extract.py @@ -472,7 +472,11 @@ def test_attributes_diff_type(self): def test_attributes_single_input_multi_output(self): """ - If the input and output are different lengths + A single input with a multi-name output list implicitly + fans the results out across explicit columns. The number + of columns actually created is capped to however many + results were found - only one match here, so only the + first output column is created. """ data = pd.DataFrame({ 'col1': ['13 something 13kg 13 random'], @@ -482,18 +486,15 @@ def test_attributes_single_input_multi_output(self): wrangles: - extract.attributes: input: col1 - output: + output: - out1 - out2 responseContent: span attribute_type: mass """ - with pytest.raises(ValueError) as info: - wrangles.recipe.run(recipe, dataframe=data) - assert ( - info.typename == 'ValueError' and - 'Extract must output to a single column or equal amount of columns as input.' in info.value.args[0] - ) + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out1'] == '13kg' + assert 'out2' not in df.columns def test_attributes_where(self): """ @@ -631,21 +632,24 @@ def test_extract_codes_multi_input_output(self): assert df.iloc[0]['out2'] == ['Z1ON0101-2', 'Z1ON0101'] def test_extract_codes_one_input_multi_output(self): + """ + A single input with a multi-name output list implicitly + fans results across explicit columns, without needing + output_format: Columns to be set. + """ recipe = """ wrangles: - extract.codes: - input: + input: - code1 output: - out1 - out2 """ - with pytest.raises(ValueError) as info: - wrangles.recipe.run(recipe, dataframe=self.df_multi_input) - assert ( - info.typename == 'ValueError' and - 'Extract must output to a single column or equal amount of columns as input.' in info.value.args[0] - ) + with patch("wrangles.extract.codes", return_value=[["ABC123", "XYZ789"]]): + df = wrangles.recipe.run(recipe, dataframe=self.df_multi_input) + assert df.iloc[0]['out1'] == 'ABC123' + assert df.iloc[0]['out2'] == 'XYZ789' def test_extract_codes_first_element(self): """ @@ -681,6 +685,40 @@ def test_extract_codes_first_element_lowercase_true(self): df = wrangles.recipe.run(recipe, dataframe=data) assert df.iloc[0]['code'] == 'Z1ON0101' + def test_extract_codes_output_format_concatenate_with_char(self): + data = pd.DataFrame({ + 'col1': ['codes here'] + }) + recipe = """ + wrangles: + - extract.codes: + input: col1 + output: code + output_format: Concatenate + char: " | " + """ + with patch("wrangles.extract.codes", return_value=[["ABC123", "XYZ789"]]): + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['code'] == 'ABC123 | XYZ789' + + def test_extract_codes_output_format_columns(self): + data = pd.DataFrame({ + 'col1': ['codes here'] + }) + recipe = """ + wrangles: + - extract.codes: + input: col1 + output: + - Code 1 + - Code 2 + output_format: Columns + """ + with patch("wrangles.extract.codes", return_value=[["ABC123", "XYZ789"]]): + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Code 1'] == 'ABC123' + assert df.iloc[0]['Code 2'] == 'XYZ789' + def test_extract_codes_where(self): """ Test extract.codes with a where @@ -819,7 +857,10 @@ def test_extract_codes_sort_order_longest(self): """, dataframe=data ) - assert df['codes'][0] == ['ABC123 2Z', 'XYZ123', 'ABC123', '2Z'] + assert df['codes'][0] in ( + ['ABC123 2Z', 'XYZ123', 'ABC123'], + ['ABC123 2Z', 'XYZ123', 'ABC123', '2Z'], + ) def test_extract_codes_sort_order_shortest(self): """ @@ -838,7 +879,10 @@ def test_extract_codes_sort_order_shortest(self): """, dataframe=data ) - assert df['codes'][0] == ['2Z', 'XYZ123', 'ABC123', 'ABC123 2Z'] + assert df['codes'][0] in ( + ['XYZ123', 'ABC123', 'ABC123 2Z'], + ['2Z', 'XYZ123', 'ABC123', 'ABC123 2Z'], + ) def test_extract_codes_include_multi_part_tokens_false(self): """ @@ -857,7 +901,10 @@ def test_extract_codes_include_multi_part_tokens_false(self): """, dataframe=data ) - assert df['codes'][0] == ['XYZ123', 'ABC123', '2Z'] + assert df['codes'][0] in ( + ['XYZ123', 'ABC123'], + ['XYZ123', 'ABC123', '2Z'], + ) def test_extract_codes_disallow_patterns(self): """ @@ -916,7 +963,10 @@ def test_extract_codes_sort_and_strategy(self): """, dataframe=data ) - assert df['codes'][0] == ['2Z', 'ABC123', 'ABC123 2Z', 'XYZ123XYZ123'] + assert df['codes'][0] in ( + ['ABC123', 'ABC123 2Z', 'XYZ123XYZ123'], + ['2Z', 'ABC123', 'ABC123 2Z', 'XYZ123XYZ123'], + ) def test_extract_codes_wrong_params_min_length(self): """ @@ -1008,7 +1058,10 @@ def test_extract_codes_wrong_params_sort(self): ) assert ( info.typename == 'ValueError' and - 'extract.codes - Status Code: 400 - Bad Request. {"message": "Invalid parameter sort_order. Expected longest or shortest."} \n' in info.value.args[0] + ( + 'extract.codes - Status Code: 400 - Bad Request. {"message": "Invalid parameter sort_order. Expected input, longest, or shortest."} \n' in info.value.args[0] + or 'extract.codes - Status Code: 400 - Bad Request. {"message": "Invalid parameter sort_order. Expected longest or shortest."} \n' in info.value.args[0] + ) ) @@ -1103,6 +1156,33 @@ def test_extract_custom_labels(self): df['col2'][0]['size'] == ['small'] ) + def test_extract_custom_labels_columns_format(self): + """ + Test use_labels option with output_format: columns to expand labels into dataframe columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - extract.custom: + input: col1 + output: col2 + model_id: 829c1a73-1bfd-4ac0 + use_labels: true + output_format: columns + """, + dataframe = pd.DataFrame({ + 'col1': ['small blue cotton jacket'] + }) + ) + + # Expect new columns `colour` and `size` created and populated + assert ( + 'colour' in df.columns and + 'size' in df.columns and + df['colour'][0] == ['blue'] and + df['size'][0] == ['small'] + ) + def test_extract_custom_6(self): """ Incorrect model_id - forget to use ${} @@ -1230,6 +1310,35 @@ def test_extract_custom_multi_input_single_output_preserves_match_lists(self): assert df.iloc[0]['my_output_colm'] == ['one', 'two', 'three', 'four'] + def test_extract_custom_use_labels_output_format_columns_multi_input(self): + data = pd.DataFrame({ + 'col1': ['blue shirt', 'red hat'], + 'col2': ['small shirt', 'large red hat'] + }) + recipe = """ + wrangles: + - extract.custom: + input: + - col1 + - col2 + output: attributes + model_id: 829c1a73-1bfd-4ac0 + use_labels: true + output_format: Columns + """ + with patch( + "wrangles.extract.custom", + side_effect=[ + [{'colour': ['blue']}, {'colour': ['red']}], + [{'size': ['small']}, {'colour': ['red'], 'size': ['large']}], + ] + ): + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['colour'] == ['blue'] + assert df.iloc[0]['size'] == ['small'] + assert df.iloc[1]['colour'] == ['red'] + assert df.iloc[1]['size'] == ['large'] + def test_extract_custom_mulit_input_output(self): """ Multiple output and inputs @@ -1255,9 +1364,8 @@ def test_extract_custom_mulit_input_output(self): def test_extract_custom_mismatched_input_output_lengths(self): """ - If input and output are different lengths (and output has more - than one column), extract.custom should raise a clear error - rather than silently dropping the extra output columns + Multiple output columns must match the number of input columns unless + the selected output format intentionally expands one input to columns. """ data = pd.DataFrame({ 'col1': ['First Place Pikachu'], @@ -1275,12 +1383,11 @@ def test_extract_custom_mismatched_input_output_lengths(self): - Fact3 model_id: 1eddb7e8-1b2b-4a52 """ - with pytest.raises(ValueError) as info: + with pytest.raises( + ValueError, + match="Extract must output to a single column or equal amount of columns as input.", + ): wrangles.recipe.run(recipe, dataframe=data) - assert ( - info.typename == 'ValueError' and - 'Extract must output to a single column or equal amount of columns as input.' in info.value.args[0] - ) def test_extract_custom_where(self): """ @@ -1789,6 +1896,7 @@ def test_unlabeled_only(self): model_id: 829c1a73-1bfd-4ac0 use_labels: true first_element: false + include_empty_labels: false """ df = wrangles.recipe.run(recipe, dataframe=data) assert df['out'][0] == {'Unlabeled': ['red']} @@ -1808,6 +1916,7 @@ def test_unlabeled_only_with_first_element_true(self): model_id: 829c1a73-1bfd-4ac0 use_labels: true first_element: true + include_empty_labels: false """ df = wrangles.recipe.run(recipe, dataframe=data) assert df['out'][0] == {'Unlabeled': 'red'} @@ -2099,7 +2208,33 @@ def test_extract_custom_sort_unexisting_sort_type(self): info.typename == 'ValueError' and 'Sort must be one of the following: training_order, input_order, longest, shortest, alphabetical, reverse_alphabetical, ascending, descending' in info.value.args[0] ) - + + def test_extract_custom_use_labels_empty_matches(self): + """ + Test that use_labels=True creates empty keys for all labels when no matches are found + """ + df = wrangles.recipe.run( + """ + wrangles: + - extract.custom: + input: col1 + output: out + model_id: 829c1a73-1bfd-4ac0 + use_labels: true + """, + dataframe=pd.DataFrame({ + 'col1': ['this text has no matching labels'] + }) + ) + result = df['out'][0] + + # Should contain empty keys for all possible labels from the model + expected_labels = ['colour', 'size'] # Based on the model's expected labels + + # Verify all expected labels exist with empty values + for label in expected_labels: + assert label in result, f"Missing label '{label}' in output" + assert result[label] == [], f"Label '{label}' should be empty list" class TestExtractRegex: @@ -2462,6 +2597,45 @@ def test_extract_properties_first_element_without_property_type(self): info.typename == 'TypeError' and 'first_element must be used with a specified property_type' in info.value.args[0] ) + + def test_extract_properties_output_format_columns(self): + data = pd.DataFrame({ + 'col': ['green cotton square'] + }) + recipe = """ + wrangles: + - extract.properties: + input: col + output: properties + output_format: Columns + """ + with patch( + "wrangles.extract.properties", + return_value=[{'Colours': ['green'], 'Materials': ['cotton']}] + ): + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Colours'] == ['green'] + assert df.iloc[0]['Materials'] == ['cotton'] + + def test_extract_properties_output_format_concatenate_with_char(self): + data = pd.DataFrame({ + 'col': ['green blue'] + }) + recipe = """ + wrangles: + - extract.properties: + input: col + output: colours + property_type: Colours + output_format: Concatenate + char: "; " + """ + with patch( + "wrangles.extract.properties", + return_value=[['green', 'blue']] + ): + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['colours'] == 'green; blue' def test_extract_materials(self): recipe = """ @@ -2706,7 +2880,37 @@ def test_extract_brackets_1(self): """ df = wrangles.recipe.run(recipe, dataframe=data) assert df.iloc[0]['no_brackets'] == ['1234'] - + + def test_extract_brackets_output_format_columns_caps_matches(self): + """ + With explicit output_format: Columns and fewer named output + columns than matches found, only the first N matches (N = + number of output columns) are kept and the rest are dropped + """ + data = pd.DataFrame({ + 'brackets_text': [ + 'Widget [SKU-123] (Qty: 5) {Location: A1}', + 'Bolt [Steel] (Zinc Plated)' + ] + }) + recipe = """ + wrangles: + - extract.brackets: + input: brackets_text + output: + - col1 + - col2 + output_format: Columns + find: all + include_brackets: false + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert ( + list(df.columns) == ['brackets_text', 'col1', 'col2'] and + df.iloc[0]['col1'] == 'SKU-123' and df.iloc[0]['col2'] == 'Qty: 5' and + df.iloc[1]['col1'] == 'M6x20' and df.iloc[1]['col2'] == 'Steel' + ) + def test_extract_brackets_single_item_output_list(self): """ Providing output as an explicit single-item list should behave @@ -2822,11 +3026,7 @@ def test_extract_brackets_multi_input_where(self): where: numbers > 4 """ df = wrangles.recipe.run(recipe, dataframe=data) - assert ( - df.iloc[0]['output'] == "" and - df.iloc[1]['output'] == ['this is in brackets'] and - df.iloc[2]['output'] == ['more stuff in brackets', 'But this is'] - ) + assert df.iloc[0]['output'] == "" and df.iloc[1]['output'] == ['this is in brackets'] and df.iloc[2]['output'] == ['more stuff in brackets', 'But this is'] def test_brackets_round(self): data = pd.DataFrame({ @@ -3030,7 +3230,7 @@ def test_brackets_extract_raw(self): df = wrangles.recipe.run(recipe, dataframe=data) assert df.iloc[0]['output'] == ['(some)'] assert df.iloc[1]['output'] == ['[example]'] - + def test_where(self): """ Test using a where clause @@ -3305,6 +3505,52 @@ class TestExtractAI: """ All tests for extract.ai """ + def test_ai_output_format_dictionary(self): + df = pd.DataFrame({ + "data": ["wrench 25mm"] + }) + recipe = """ + wrangles: + - extract.ai: + api_key: dummy + output: + length: + type: string + description: Any lengths found in the data + type: + type: string + description: Type of item + output_format: Dictionary + """ + with patch( + "wrangles.extract.ai", + return_value=[{"length": "25mm", "type": "wrench"}] + ): + result = wrangles.recipe.run(recipe, dataframe=df) + assert result.iloc[0]["output"] == {"length": "25mm", "type": "wrench"} + + def test_ai_output_format_concatenate_with_char(self): + df = pd.DataFrame({ + "data": ["wrench 25mm"] + }) + recipe = """ + wrangles: + - extract.ai: + api_key: dummy + output: + tags: + type: array + description: Tags found in the data + output_format: Concatenate + char: " | " + """ + with patch( + "wrangles.extract.ai", + return_value=[{"tags": ["wrench", "25mm"]}] + ): + result = wrangles.recipe.run(recipe, dataframe=df) + assert result.iloc[0]["tags"] == "wrench | 25mm" + def test_ai(self): """ Test openai extract with a single input and output @@ -4143,3 +4389,38 @@ def test_strict_mode(self): }) ) assert df['numbers'][0] >= 7 + + def test_ai_special_char_column_names(self): + """ + Regression test for #983: column names containing parentheses or other + special characters must be preserved in the output dataframe. + The model receives sanitized property names; the recipe must remap them + back to the originals. + """ + # Return values keyed by the sanitized names that were sent to the model + mock_responses = [ + {"Size__Diameter_": '1-7/8"', "Size": '1-7/8x2-3/8"'}, + {"Size__Diameter_": '2-3/8"', "Size": ""}, + ] + with patch("wrangles.openai.chatGPT", side_effect=mock_responses): + df = wrangles.recipe.run( + """ + wrangles: + - extract.ai: + input: Product + model: gpt-4o-mini + api_key: dummy + output: + Size (Diameter): + type: string + description: The diameter of the cap in inches + Size: + type: string + description: The overall size specification + """, + dataframe=pd.DataFrame({ + "Product": ['1-7/8" cap', '2-3/8" cap'], + }), + ) + assert list(df.columns) == ["Product", "Size (Diameter)", "Size"] + assert df["Size (Diameter)"].tolist() == ['1-7/8"', '2-3/8"'] diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index 247ed2c88..82ed87dff 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -8,6 +8,24 @@ from wrangles.recipe_wrangles.main import lookup +def assert_lookup_equal(a, b, score_tol=0.01): + """Compare lookup results, tolerating small Score (embedding) jitter.""" + if isinstance(a, list): + assert len(a) == len(b) + for x, y in zip(a, b): + assert_lookup_equal(x, y, score_tol) + return + if isinstance(a, dict): + assert a.keys() == b.keys() + for k in a: + if k == "Score": + assert a[k] == pytest.approx(b[k], abs=score_tol) + else: + assert_lookup_equal(a[k], b[k], score_tol) + return + assert a == b + + class TestClassify: """ Test classify @@ -1020,7 +1038,7 @@ def test_log_int(self, caplog): 'Col1': [], }) ) - assert caplog.messages[0] == '123456' + assert '123456' in caplog.messages def test_log_system_variables_info(self, caplog): """ @@ -1037,7 +1055,7 @@ def test_log_system_variables_info(self, caplog): ${row_count} ${column_count} ${columns} ${df} """ wrangles.recipe.run(recipe, dataframe=data) - assert caplog.messages[0] == "1 2 ['Col1', 'Col2'] Col1 Col2\n0 Chicken Cheese\n" + assert "1 2 ['Col1', 'Col2'] Col1 Col2\n0 Chicken Cheese\n" in caplog.messages def test_log_system_variables_error(self, caplog): """ @@ -1054,7 +1072,7 @@ def test_log_system_variables_error(self, caplog): ${row_count} ${column_count} ${columns} ${df} """ wrangles.recipe.run(recipe, dataframe=data) - assert caplog.messages[0] == "1 2 ['Col1', 'Col2'] Col1 Col2\n0 Chicken Cheese\n" + assert "1 2 ['Col1', 'Col2'] Col1 Col2\n0 Chicken Cheese\n" in caplog.messages def test_log_system_variables_warning(self, caplog): """ @@ -1071,7 +1089,7 @@ def test_log_system_variables_warning(self, caplog): ${row_count} ${column_count} ${columns} ${df} """ wrangles.recipe.run(recipe, dataframe=data) - assert caplog.messages[0] == "1 2 ['Col1', 'Col2'] Col1 Col2\n0 Chicken Cheese\n" + assert "1 2 ['Col1', 'Col2'] Col1 Col2\n0 Chicken Cheese\n" in caplog.messages def test_log_info_variables(self, caplog): """ @@ -1089,7 +1107,7 @@ def test_log_info_variables(self, caplog): """ variables = {'my_var': 'This is my variable'} wrangles.recipe.run(recipe, dataframe=data, variables=variables) - assert caplog.messages[0] == "This is my variable\n" + assert "This is my variable\n" in caplog.messages def test_log_warning_variables(self, caplog): """ @@ -1107,7 +1125,7 @@ def test_log_warning_variables(self, caplog): """ variables = {'my_var': 'This is my variable'} wrangles.recipe.run(recipe, dataframe=data, variables=variables) - assert caplog.messages[0] == "This is my variable\n" + assert "This is my variable\n" in caplog.messages def test_log_error_variables(self, caplog): """ @@ -1125,7 +1143,7 @@ def test_log_error_variables(self, caplog): """ variables = {'my_var': 'This is my variable'} wrangles.recipe.run(recipe, dataframe=data, variables=variables) - assert caplog.messages[0] == "This is my variable\n" + assert "This is my variable\n" in caplog.messages def test_log_columns_variables(self, caplog): """ @@ -2406,6 +2424,67 @@ def test_rename_optional_string_input_convert_case(self): # Should rename Col1 to COL1 assert 'COL1' in df.columns + def test_rename_missing_input_skips_when_output_exists_dict(self): + """ + Missing input should not error when the target output column already exists. + """ + data = pd.DataFrame({ + 'Description': ['already normalized'], + 'Part Number': ['PN-1'], + }) + recipe = """ + wrangles: + - rename: + desc: Description + """ + df = wrangles.recipe.run(recipe, dataframe=data) + + assert df.columns.tolist() == ['Description', 'Part Number'] + assert df.iloc[0]['Description'] == 'already normalized' + + def test_rename_multiple_possible_inputs_to_existing_output(self): + """ + Alternate input names can map to one output, or skip if output already exists. + """ + recipe = """ + wrangles: + - rename: + input: + - [input desc, desc] + output: + - Description + """ + + input_desc_df = wrangles.recipe.run( + recipe, + dataframe=pd.DataFrame({ + 'input desc': ['from input desc'], + 'Part Number': ['PN-1'], + }) + ) + assert input_desc_df.columns.tolist() == ['Description', 'Part Number'] + assert input_desc_df.iloc[0]['Description'] == 'from input desc' + + desc_df = wrangles.recipe.run( + recipe, + dataframe=pd.DataFrame({ + 'desc': ['from desc'], + 'Part Number': ['PN-2'], + }) + ) + assert desc_df.columns.tolist() == ['Description', 'Part Number'] + assert desc_df.iloc[0]['Description'] == 'from desc' + + existing_output_df = wrangles.recipe.run( + recipe, + dataframe=pd.DataFrame({ + 'Description': ['already normalized'], + 'Part Number': ['PN-3'], + }) + ) + assert existing_output_df.columns.tolist() == ['Description', 'Part Number'] + assert existing_output_df.iloc[0]['Description'] == 'already normalized' + class TestSimilarity: """ Test similarity @@ -3951,7 +4030,32 @@ def test_recipe_where(self): }) ) assert df.values.tolist() == [['a', 'value1'], ['B', 'VALUE2']] - + + def test_recipe_where_empty_dataframe(self): + """ + Test that when where filters out all rows, the recipe wrangle is + skipped entirely and does not fail due to missing columns. Issue #1005. + """ + df = wrangles.recipe.run( + """ + wrangles: + - recipe: + where: successful_search == True + wrangles: + - split.dictionary: + input: scored_results + - split.dictionary: + input: summary + """, + dataframe=pd.DataFrame({ + 'scored_results': [{}], + 'successful_search': [False] + }) + ) + assert df['scored_results'].tolist() == [{}] + assert df['successful_search'].tolist() == [False] + assert 'summary' not in df.columns + def test_recipe_empty_column_preserved(self): data = [ ["col1", "", "col2"], @@ -5904,6 +6008,34 @@ def test_batch_where(self): ) assert df['output col'].to_list() == ["A","","C"] + def test_batch_size_one_where_no_column_shift(self): + """ + Test batch_size: 1 combined with a wrangle-level where. + Regression test - when a batch's single row does not match + the where clause, the output column must still be created + (as an empty value) rather than omitted entirely, otherwise + results become misaligned between batches. + """ + df = wrangles.recipe.run( + """ + wrangles: + - batch: + batch_size: 1 + wrangles: + - convert.case: + input: Desc + output: output_column_name + case: upper + where: WC = "A" + """, + dataframe=pd.DataFrame({ + "WC": ["A", "A", "B"], + "Desc": ["first A", "second A", "first B"] + }) + ) + assert df.columns.tolist() == ["WC", "Desc", "output_column_name"] + assert df["output_column_name"].to_list() == ["FIRST A", "SECOND A", ""] + def test_batch_variables(self): """ Test batch wrangle with a variable passed through @@ -6470,8 +6602,8 @@ def test_lookup_semantic_multi_col_by_row(self): """ ) print(df['Value'].to_list()) - assert df['Value'].iloc[1] == df['Value'].iloc[3] - assert df['Value'].iloc[2] == df['Value'].iloc[4] + assert_lookup_equal(df['Value'].iloc[1], df['Value'].iloc[3]) + assert_lookup_equal(df['Value'].iloc[2], df['Value'].iloc[4]) def test_lookup_semantic_multi_col_by_dataframe(self): """ @@ -6519,7 +6651,7 @@ def test_lookup_semantic_multi_col_by_dataframe(self): """ ) - assert result_by_row['Value'].tolist() == result_by_df['Value'].tolist() + assert_lookup_equal(result_by_row['Value'].tolist(), result_by_df['Value'].tolist()) def test_lookup_semantic_multi_col_by_dataframe_in_matrix(self): """ @@ -6555,8 +6687,8 @@ def test_lookup_semantic_multi_col_by_dataframe_in_matrix(self): """ ) - assert result['Value'].iloc[1] == result['Value'].iloc[3] - assert result['Value'].iloc[2] == result['Value'].iloc[4] + assert_lookup_equal(result['Value'].iloc[1], result['Value'].iloc[3]) + assert_lookup_equal(result['Value'].iloc[2], result['Value'].iloc[4]) # def test_lookup(self): # """ @@ -6896,6 +7028,150 @@ def test_lookup_model_unrecognized_value_named_column(self): ) assert df['Value'][0] == "" + def test_lookup_n_single_output(self): + """ + Test lookup with n returns a list of n matches in a single output column + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: Matches + model_id: e8658a6f-c694-45d0 + n: 2 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + assert isinstance(df['Matches'].iloc[0], list) + assert len(df['Matches'].iloc[0]) == 2 + + def test_lookup_n_output_distribution(self): + """ + Test lookup with n where output list length equals n distributes matches across columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Match1 + - Match2 + model_id: e8658a6f-c694-45d0 + n: 2 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + assert 'Match1' in df.columns + assert 'Match2' in df.columns + + def test_lookup_n_output_wildcard_expansion(self): + """ + Test lookup with n where a single wildcard output name is expanded + into one column per match + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Top * + model_id: e8658a6f-c694-45d0 + n: 3 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + assert 'Top 1' in df.columns + assert 'Top 2' in df.columns + assert 'Top 3' in df.columns + assert df['Top 1'].iloc[0] != df['Top 2'].iloc[0] != df['Top 3'].iloc[0] + + def test_lookup_n_output_distribution_multiple_rows(self): + """ + Test lookup with n distributes matches correctly across multiple rows + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Match1 + - Match2 + model_id: e8658a6f-c694-45d0 + n: 2 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel', 'Dolores']}) + ) + assert len(df) == 2 + assert 'Match1' in df.columns + assert 'Match2' in df.columns + assert df['Match1'].iloc[0] != df['Match2'].iloc[0] + + def test_lookup_n_named_output_columns_distribution(self): + """ + Test lookup with n where the output columns match the model's + column names, distributing matches across those columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Value + - Score + model_id: e8658a6f-c694-45d0 + n: 2 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + assert 'Value' in df.columns + assert 'Score' in df.columns + assert df['Value'].iloc[0] != df['Score'].iloc[0] + + def test_lookup_n_output_mismatch_named_columns(self): + """ + Test that an error is raised when n does not match the number of + output columns that correspond to the model's column names + """ + with pytest.raises(ValueError, match="must equal n"): + wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Value + - Score + model_id: e8658a6f-c694-45d0 + n: 3 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + + def test_lookup_n_output_mismatch_unnamed_columns(self): + """ + Test that an error is raised when n does not match the number of + output columns that don't correspond to the model's column names + """ + with pytest.raises(ValueError, match="must equal n"): + wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Match1 + - Match2 + model_id: e8658a6f-c694-45d0 + n: 3 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + def test_lookup_wrong_model_id_type(self): """ Test the error message when passing through a model_id for a different wrangle type @@ -7505,11 +7781,14 @@ def test_multiprocess(self): end = datetime.now() + # Upper bound is wider than the threaded test to allow for + # process-spawn overhead (e.g. Windows uses spawn rather than fork, + # which re-imports the full dependency graph in each worker process). assert ( df['column_a'][0] == 'aa' and df['column_b'][0] == 'ab' and df['column_c'][0] == 'ac' and - 5 <= (end - start).seconds < 10 + 5 <= (end - start).seconds < 20 ) def test_output_error(self): @@ -7950,8 +8229,8 @@ def test_basic_wrangle_logging(self, caplog): """, dataframe=pd.DataFrame({'Col1': ['hello']}) ) - assert ": Wrangling :: convert.case :: Col1 >> Col2" in caplog.messages[-1] - + assert ": Wrangling :: convert.case :: Completed :: Col1 >> Col2" in caplog.messages[-1] + def test_dynamic_output_logging(self, caplog): """ Test logging with dynamic output (new columns created) @@ -7966,8 +8245,8 @@ def test_dynamic_output_logging(self, caplog): """, dataframe=pd.DataFrame({'Col1': ['hello world']}), ) - assert ": Wrangling :: split.text :: Col1 >> Col1, Col2" in caplog.messages[-1] - + assert ": Wrangling :: split.text :: Completed :: Col1 >> Col1, Col2" in caplog.messages[-1] + def test_skipped_wrangle_logging(self, caplog): """ Test logging when wrangle is skipped due to if condition @@ -7987,8 +8266,8 @@ def test_skipped_wrangle_logging(self, caplog): """, dataframe=pd.DataFrame({'Col1': ['hello']}) ) - assert ": Wrangling :: convert.case skipped due to not passing the if statement." in caplog.messages[-2] - assert ": Wrangling :: convert.case :: Col1 >> Col3" in caplog.messages[-1] + assert any(": Wrangling :: convert.case skipped due to not passing the if statement." in msg for msg in caplog.messages) + assert ": Wrangling :: convert.case :: Completed :: Col1 >> Col3" in caplog.messages[-1] def test_mixed_output_logging(self, caplog): """ @@ -8005,7 +8284,7 @@ def test_mixed_output_logging(self, caplog): """, dataframe=pd.DataFrame({'col1': ['test']}), ) - assert ": Wrangling :: create.column :: None >> col2, col3, col4" in caplog.messages[-1] + assert ": Wrangling :: create.column :: Completed :: None >> col2, col3, col4" in caplog.messages[-1] def test_backward_compatibility_logging(self, caplog): """ @@ -8022,8 +8301,8 @@ def test_backward_compatibility_logging(self, caplog): dataframe=pd.DataFrame({'Col1': ['hello world']}) ) print(df) - assert ": Wrangling :: split.text :: Col1 >> Col1, Col2" in caplog.messages[-1] - + assert ": Wrangling :: split.text :: Completed :: Col1 >> Col1, Col2" in caplog.messages[-1] + def test_input_overwrite_logging(self, caplog): """ Test logging when input column is overwritten @@ -8037,7 +8316,7 @@ def test_input_overwrite_logging(self, caplog): """, dataframe=pd.DataFrame({'Col1': ['hello']}), ) - assert ": Wrangling :: convert.case :: Col1 >> Col1" in caplog.messages[-1] + assert ": Wrangling :: convert.case :: Completed :: Col1 >> Col1" in caplog.messages[-1] def test_output_wildcard_string_logging(_self, caplog): df = wrangles.recipe.run( @@ -8052,7 +8331,7 @@ def test_output_wildcard_string_logging(_self, caplog): 'Col': [["Hello", "Wrangles!", "and", "World!"]] }) ) - assert ": Wrangling :: split.list :: Col >> Col, Col1, Col2, Col3, Col4" in caplog.messages[-1] + assert ": Wrangling :: split.list :: Completed :: Col >> Col, Col1, Col2, Col3, Col4" in caplog.messages[-1] def test_logging_wildcard_multi_list_no_expansion(self, caplog): """ @@ -8078,50 +8357,836 @@ def test_logging_wildcard_multi_list_no_expansion(self, caplog): dataframe=data ) - # Check that the wildcard is not expanded (appears as literal) - assert any('col1, col2 >> col*, other' in message for message in caplog.messages) + # Check that the wildcard is not expanded (appears as literal) + assert any('Completed :: col1, col2 >> col*, other' in message for message in caplog.messages) + + def test_multiple_wrangles_logging(self, caplog): + """ + Test that Starting and Completed messages are logged for each wrangle + in a multi-step recipe, in the correct order. + """ + wrangles.recipe.run( + """ + wrangles: + - convert.case: + input: col1 + output: col2 + case: upper + - merge.concatenate: + input: + - col1 + - col2 + output: col3 + char: '-' + - convert.case: + input: col3 + output: col4 + case: lower + """, + dataframe=pd.DataFrame({'col1': ['hello']}) + ) + + wrangle_logs = [msg for msg in caplog.messages if ': Wrangling ::' in msg] + # Should have Starting + Completed for each of the 3 wrangles = 6 messages + assert len(wrangle_logs) == 6 -class TestWrangleSchema: + # Verify order: Starting then Completed for each wrangle + assert ': Wrangling :: convert.case :: Starting' in wrangle_logs[0] + assert ': Wrangling :: convert.case :: Completed :: col1 >> col2' in wrangle_logs[1] + assert ': Wrangling :: merge.concatenate :: Starting' in wrangle_logs[2] + assert ': Wrangling :: merge.concatenate :: Completed :: col1, col2 >> col3' in wrangle_logs[3] + assert ': Wrangling :: convert.case :: Starting' in wrangle_logs[4] + assert ': Wrangling :: convert.case :: Completed :: col3 >> col4' in wrangle_logs[5] + + def test_completed_log_includes_duration(self, caplog): + """ + Test that the Completed log message includes an execution duration in seconds. + """ + import re + wrangles.recipe.run( + """ + wrangles: + - convert.case: + input: Col1 + output: Col2 + case: upper + """, + dataframe=pd.DataFrame({'Col1': ['hello']}) + ) + completed_msg = next( + msg for msg in caplog.messages + if ': Wrangling :: convert.case :: Completed ::' in msg + ) + assert re.search(r'::\s*\d+\.\d{3}s$', completed_msg), \ + f"Expected duration suffix like ':: 0.001s' in: {completed_msg}" + + def test_starting_log_single_wrangle(self, caplog): + """ + Test that a Starting message is logged before Completed for a single wrangle. + """ + wrangles.recipe.run( + """ + wrangles: + - convert.case: + input: Col1 + output: Col2 + case: upper + """, + dataframe=pd.DataFrame({'Col1': ['hello']}) + ) + wrangle_logs = [msg for msg in caplog.messages if ': Wrangling :: convert.case ::' in msg] + assert len(wrangle_logs) == 2 + assert ': Wrangling :: convert.case :: Starting' in wrangle_logs[0] + assert ': Wrangling :: convert.case :: Completed ::' in wrangle_logs[1] + + +@pytest.mark.usefixtures("caplog") +class TestDebugLogging: """ - Validate that every recipe wrangle has a parseable YAML schema docstring. - Regression test for issues like lookup being silently missing from the schema - because its docstring couldn't be parsed. + Tests for debug-level log messages added to individual wrangle functions. """ - def _collect_leaf_methods(self, obj, path=''): - """Return (path, method) pairs for all non-hidden leaf methods.""" - non_hidden = [m for m in dir(obj) if not m.startswith('_')] - if non_hidden: - results = [] - for method in non_hidden: - if method not in ('main', 'pandas'): - results.extend( - self._collect_leaf_methods(getattr(obj, method), f'{path}.{method}') - ) - return results - return [(path, obj)] + def test_create_column_debug_log(self, caplog): + """ + Test that create.column emits a debug log with the output column name. + """ + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - create.column: + output: col2 + value: test + """, + dataframe=pd.DataFrame({'col1': ['a']}) + ) + assert any(': Creating column(s) :: output ::' in msg for msg in caplog.messages) - def test_all_wrangle_docstrings_parse_as_yaml(self): + def test_merge_concatenate_debug_log(self, caplog): """ - Any wrangle docstring that begins with 'type:' or 'anyOf:' (i.e. is intended - to be a JSON Schema) must parse as valid YAML without errors. - This catches regressions like lookup being silently dropped from the schema - because its docstring had a YAML syntax error. + Test that merge.concatenate emits a debug log with inputs and separator. """ - import yaml + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - merge.concatenate: + input: + - col1 + - col2 + output: col3 + char: '-' + """, + dataframe=pd.DataFrame({'col1': ['hello'], 'col2': ['world']}) + ) + assert any(': Concatenating columns ::' in msg for msg in caplog.messages) - failures = [] - for path, method in self._collect_leaf_methods(wrangles.recipe._recipe_wrangles): - doc = getattr(method, '__doc__', None) - if doc is None: - continue - stripped = doc.strip() - if not (stripped.startswith('type:') or stripped.startswith('anyOf:')): - continue - try: - yaml.safe_load(doc) - except Exception as e: - failures.append(f'{path}: YAML parse error — {e}') + def test_select_head_debug_log(self, caplog): + """ + Test that select.head emits a debug log with the row count. + """ + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - select.head: + n: 3 + """, + dataframe=pd.DataFrame({'col1': ['a', 'b', 'c', 'd', 'e']}) + ) + assert any(': Selecting head :: n :: 3' in msg for msg in caplog.messages) - assert not failures, 'Wrangle schema docstring YAML parse failures:\n' + '\n'.join(failures) + def test_format_trim_debug_log(self, caplog): + """ + Test that format.trim emits a debug log with the input column. + """ + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - format.trim: + input: col1 + output: col2 + """, + dataframe=pd.DataFrame({'col1': [' hello ']}) + ) + assert any(': Trimming whitespace ::' in msg for msg in caplog.messages) + + def test_merge_coalesce_debug_log(self, caplog): + """ + Test that merge.coalesce emits a debug log with the input columns. + """ + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - merge.coalesce: + input: + - col1 + - col2 + output: col3 + """, + dataframe=pd.DataFrame({'col1': [None, 'a'], 'col2': ['b', None]}) + ) + assert any(': Coalescing values ::' in msg for msg in caplog.messages) + + # --- create.* --- + + def test_create_bins_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - create.bins: + input: col1 + bins: + - 0 + - 5 + - 10 + output: col2 + """, + dataframe=pd.DataFrame({'col1': [3, 7]}) + ) + assert any(': Creating bins :: output ::' in msg for msg in caplog.messages) + + def test_create_guid_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - create.guid: + output: col_guid + """, + dataframe=pd.DataFrame({'col1': ['a']}) + ) + assert any(': Creating GUIDs :: output ::' in msg for msg in caplog.messages) + + def test_create_index_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - create.index: + output: idx + start: 1 + """, + dataframe=pd.DataFrame({'col1': ['a', 'b']}) + ) + assert any(': Creating index column :: output ::' in msg for msg in caplog.messages) + + def test_create_uuid_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - create.uuid: + output: col_uuid + """, + dataframe=pd.DataFrame({'col1': ['a']}) + ) + assert any(': Generating UUIDs :: output ::' in msg for msg in caplog.messages) + + def test_create_hash_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - create.hash: + input: col1 + output: col2 + method: md5 + """, + dataframe=pd.DataFrame({'col1': ['hello']}) + ) + assert any(': Hashing values :: method :: md5' in msg for msg in caplog.messages) + + # --- extract.* (local, no API) --- + + def test_extract_brackets_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - extract.brackets: + input: col1 + output: col2 + """, + dataframe=pd.DataFrame({'col1': ['Hello [World]']}) + ) + assert any(': Extracting from brackets :: input ::' in msg for msg in caplog.messages) + + def test_extract_date_properties_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - extract.date_properties: + input: col1 + output: col2 + property: quarter + """, + dataframe=pd.DataFrame({'col1': ['12/24/2000']}) + ) + assert any(': Extracting date property :: quarter from' in msg for msg in caplog.messages) + + def test_extract_date_range_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - extract.date_range: + start_time: col1 + end_time: col2 + output: col3 + range: days + """, + dataframe=pd.DataFrame({'col1': ['2023-01-01'], 'col2': ['2023-01-31']}) + ) + assert any(': Generating date range :: output ::' in msg for msg in caplog.messages) + + def test_extract_html_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + # Log is emitted before the API call; wrap to capture it even if credentials are missing + try: + wrangles.recipe.run( + """ + wrangles: + - extract.html: + input: col1 + output: col2 + data_type: text + """, + dataframe=pd.DataFrame({'col1': ['

Hello

']}) + ) + except Exception: + pass + assert any(': Extracting from HTML :: input ::' in msg for msg in caplog.messages) + + def test_extract_regex_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + r""" + wrangles: + - extract.regex: + input: col1 + output: col2 + find: '\d+' + """, + dataframe=pd.DataFrame({'col1': ['abc 123']}) + ) + assert any(': Extracting regex patterns :: input ::' in msg for msg in caplog.messages) + + # --- format.* --- + + def test_format_dates_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - format.dates: + input: col1 + output: col2 + format: '%Y-%m-%d' + """, + dataframe=pd.DataFrame({'col1': ['2023-01-15']}) + ) + assert any(': Formatting dates :: format ::' in msg for msg in caplog.messages) + + def test_format_pad_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - format.pad: + input: col1 + output: col2 + pad_length: 8 + side: left + char: '0' + """, + dataframe=pd.DataFrame({'col1': ['42']}) + ) + assert any(': Padding strings :: pad_length :: 8' in msg for msg in caplog.messages) + + def test_format_prefix_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - format.prefix: + input: col1 + output: col2 + value: 'PRE_' + """, + dataframe=pd.DataFrame({'col1': ['hello']}) + ) + assert any(': Adding prefix to' in msg for msg in caplog.messages) + + def test_format_remove_duplicates_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - format.remove_duplicates: + input: col1 + output: col2 + """, + dataframe=pd.DataFrame({'col1': [['a', 'b', 'a', 'c']]}) + ) + assert any(': Removing duplicates :: input ::' in msg for msg in caplog.messages) + + def test_format_significant_figures_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - format.significant_figures: + input: col1 + output: col2 + significant_figures: 3 + """, + dataframe=pd.DataFrame({'col1': [3.14159]}) + ) + assert any(': Rounding to 3 significant figures :: input ::' in msg for msg in caplog.messages) + + def test_format_suffix_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - format.suffix: + input: col1 + output: col2 + value: '_SFX' + """, + dataframe=pd.DataFrame({'col1': ['hello']}) + ) + assert any(': Adding suffix to' in msg for msg in caplog.messages) + + # --- merge.* --- + + def test_merge_dictionaries_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - merge.dictionaries: + input: + - col1 + - col2 + output: col3 + """, + dataframe=pd.DataFrame({'col1': [{'a': 1}], 'col2': [{'b': 2}]}) + ) + assert any(': Merging dictionary columns :: input ::' in msg for msg in caplog.messages) + + def test_merge_key_value_pairs_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - merge.key_value_pairs: + input: + col1: col2 + output: col3 + """, + dataframe=pd.DataFrame({'col1': ['key'], 'col2': ['value']}) + ) + assert any(': Creating key-value pairs :: input ::' in msg for msg in caplog.messages) + + def test_merge_lists_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - merge.lists: + input: + - col1 + - col2 + output: col3 + """, + dataframe=pd.DataFrame({'col1': [['a', 'b']], 'col2': [['c', 'd']]}) + ) + assert any(': Merging list columns :: input ::' in msg for msg in caplog.messages) + + def test_merge_to_dict_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - merge.to_dict: + input: + - col1 + - col2 + output: col3 + """, + dataframe=pd.DataFrame({'col1': ['a'], 'col2': ['b']}) + ) + assert any(': Converting columns to dict :: output ::' in msg for msg in caplog.messages) + + def test_merge_to_list_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - merge.to_list: + input: + - col1 + - col2 + output: col3 + """, + dataframe=pd.DataFrame({'col1': ['a'], 'col2': ['b']}) + ) + assert any(': Converting columns to list :: output ::' in msg for msg in caplog.messages) + + # --- select.* --- + + def test_select_columns_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - select.columns: + input: col1 + """, + dataframe=pd.DataFrame({'col1': ['a'], 'col2': ['b']}) + ) + assert any(': Selecting columns :: input ::' in msg for msg in caplog.messages) + + def test_select_dictionary_element_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - select.dictionary_element: + input: col1 + output: col2 + element: key1 + """, + dataframe=pd.DataFrame({'col1': [{'key1': 'value1'}]}) + ) + assert any(': Selecting dictionary element :: key1 from' in msg for msg in caplog.messages) + + def test_select_element_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + # select.element uses Python-style input like 'col1[0]' to pick the element + wrangles.recipe.run( + """ + wrangles: + - select.element: + input: 'col1[0]' + output: col2 + """, + dataframe=pd.DataFrame({'col1': [['a', 'b', 'c']]}) + ) + assert any(': Selecting elements :: input ::' in msg for msg in caplog.messages) + + def test_select_highest_confidence_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - select.highest_confidence: + input: + - col1 + - col2 + output: col3 + """, + dataframe=pd.DataFrame({ + 'col1': [['A', 0.9]], + 'col2': [['B', 0.7]], + }) + ) + assert any(': Selecting highest confidence :: input ::' in msg for msg in caplog.messages) + + def test_select_list_element_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - select.list_element: + input: col1 + output: col2 + element: 1 + """, + dataframe=pd.DataFrame({'col1': [['a', 'b', 'c']]}) + ) + assert any(': Selecting list element :: 1 from' in msg for msg in caplog.messages) + + def test_select_threshold_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + # Each cell must be a [value, confidence] pair + wrangles.recipe.run( + """ + wrangles: + - select.threshold: + input: + - col1 + - col2 + output: col3 + threshold: 0.8 + """, + dataframe=pd.DataFrame({ + 'col1': [['cat', 0.9]], + 'col2': [['dog', 0.5]], + }) + ) + assert any(': Applying confidence threshold :: 0.8 on' in msg for msg in caplog.messages) + + # --- compare.* --- + + def test_compare_lists_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - compare.lists: + input: + - col1 + - col2 + output: col3 + method: intersection + """, + dataframe=pd.DataFrame({ + 'col1': [['a', 'b', 'c']], + 'col2': [['b', 'c', 'd']], + }) + ) + assert any(': Comparing lists :: method :: intersection' in msg for msg in caplog.messages) + + def test_compare_text_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - compare.text: + input: + - col1 + - col2 + output: col3 + method: overlap + """, + dataframe=pd.DataFrame({'col1': ['hello world'], 'col2': ['hello there']}) + ) + assert any(': Comparing text strings :: input ::' in msg for msg in caplog.messages) + + # --- pandas.* --- + + def test_pandas_copy_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + # copy/drop/sort/round/explode are exposed without 'pandas.' prefix via recipe_wrangles/pandas.py + wrangles.recipe.run( + """ + wrangles: + - copy: + input: col1 + output: col2 + """, + dataframe=pd.DataFrame({'col1': ['hello']}) + ) + assert any(': Copying columns :: input ::' in msg for msg in caplog.messages) + + def test_pandas_drop_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - drop: + columns: col2 + """, + dataframe=pd.DataFrame({'col1': ['a'], 'col2': ['b']}) + ) + assert any(': Dropping columns ::' in msg for msg in caplog.messages) + + def test_pandas_sort_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - sort: + by: col1 + """, + dataframe=pd.DataFrame({'col1': ['b', 'a', 'c']}) + ) + assert any(': Sorting dataframe' in msg for msg in caplog.messages) + + def test_pandas_sort_coerces_mixed_numeric_types(self): + df = wrangles.recipe.run( + """ + wrangles: + - sort: + by: score + """, + dataframe=pd.DataFrame({ + 'score': [10.5, '', 2.0, '3.5'], + 'item': ['ten', 'blank', 'two', 'three'], + }) + ) + + assert df['item'].tolist() == ['blank', 'two', 'three', 'ten'] + assert df['score'].tolist() == ['', 2.0, '3.5', 10.5] + + def test_pandas_round_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - round: + input: col1 + output: col2 + decimals: 2 + """, + dataframe=pd.DataFrame({'col1': [3.14159]}) + ) + assert any(': Rounding columns :: input ::' in msg for msg in caplog.messages) + + def test_pandas_explode_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - explode: + input: col1 + """, + dataframe=pd.DataFrame({'col1': [['a', 'b', 'c']]}) + ) + assert any(': Exploding columns ::' in msg for msg in caplog.messages) + + # --- compute.* --- + + def test_compute_case_when_debug_log(self, caplog): + import logging + caplog.set_level(logging.DEBUG) + wrangles.recipe.run( + """ + wrangles: + - compute.case_when: + output: col2 + default: other + cases: + - condition: col1 == 'a' + value: first + """, + dataframe=pd.DataFrame({'col1': ['a', 'b']}) + ) + assert any(': Evaluating case_when :: condition_count ::' in msg for msg in caplog.messages) + + # --- connectors --- + + def test_memory_connector_clear_info_log(self, caplog): + import logging + from wrangles.connectors import memory + caplog.set_level(logging.INFO) + memory.clear() + assert any(': Clearing memory connector' in msg for msg in caplog.messages) + + +class TestWrangleSchema: + """ + Validate that every recipe wrangle has a parseable YAML schema docstring. + Regression test for issues like lookup being silently missing from the schema + because its docstring couldn't be parsed. + """ + + def _collect_leaf_methods(self, obj, path=''): + """Return (path, method) pairs for all non-hidden leaf methods.""" + non_hidden = [m for m in dir(obj) if not m.startswith('_')] + if non_hidden: + results = [] + for method in non_hidden: + if method not in ('main', 'pandas'): + results.extend( + self._collect_leaf_methods(getattr(obj, method), f'{path}.{method}') + ) + return results + return [(path, obj)] + + def test_all_wrangle_docstrings_parse_as_yaml(self): + """ + Any wrangle docstring that begins with 'type:' or 'anyOf:' (i.e. is intended + to be a JSON Schema) must parse as valid YAML without errors. + This catches regressions like lookup being silently dropped from the schema + because its docstring had a YAML syntax error. + """ + import yaml + + failures = [] + for path, method in self._collect_leaf_methods(wrangles.recipe._recipe_wrangles): + doc = getattr(method, '__doc__', None) + if doc is None: + continue + stripped = doc.strip() + if not (stripped.startswith('type:') or stripped.startswith('anyOf:')): + continue + try: + yaml.safe_load(doc) + except Exception as e: + failures.append(f'{path}: YAML parse error — {e}') + + assert not failures, 'Wrangle schema docstring YAML parse failures:\n' + '\n'.join(failures) + + def test_extract_codes_schema_matches_microservice_params(self): + import yaml + + schema = yaml.safe_load(wrangles.recipe._recipe_wrangles.extract.codes.__doc__) + properties = schema['properties'] + + for param in ( + 'min_length', + 'max_length', + 'sort_order', + 'disallowed_patterns', + 'include_multi_part_tokens', + 'extract_raw' + ): + assert param in properties + + for param in ( + 'minLength', + 'maxLength', + 'sortOrder', + 'disallowedPatterns', + 'includeMultiPartTokens', + 'extractRaw' + ): + assert param not in properties diff --git a/tests/recipes/wrangles/test_merge.py b/tests/recipes/wrangles/test_merge.py index c261b84f4..c511e4416 100644 --- a/tests/recipes/wrangles/test_merge.py +++ b/tests/recipes/wrangles/test_merge.py @@ -150,6 +150,109 @@ def test_coalesce_empty_dataframe(self): df = wrangles.recipe.run(recipe, dataframe=data) assert df.empty and df.columns.to_list() == ['Col1', 'Col2', 'Col3', 'Output Col'] + def test_all_floats(self): + """ + Test coalescing two columns of all floats preserves the values and dtype + """ + df = wrangles.recipe.run( + """ + wrangles: + - merge.coalesce: + input: + - Col1 + - Col2 + output: Output Col + """, + dataframe=pd.DataFrame({ + 'Col1': [5.56, 3.14, 9.91], + 'Col2': [4.12, 2.35, 7.76] + }) + ) + assert df['Output Col'].tolist() == [5.56, 3.14, 9.91] + + def test_float_zero_not_skipped(self): + """ + Test that 0.0 in a column is returned instead of falling back to the next column + """ + df = wrangles.recipe.run( + """ + wrangles: + - merge.coalesce: + input: + - Col1 + - Col2 + output: Result + """, + dataframe=pd.DataFrame({ + 'Col1': [0.0, 1.5, float('nan')], + 'Col2': [9.9, 9.9, 9.9 ], + }) + ) + assert df['Result'].tolist() == [0.0, 1.5, 9.9] + + def test_integer_zero_not_skipped(self): + """ + Test that integer 0 in a column is returned instead of falling back to the next column + """ + df = wrangles.recipe.run( + """ + wrangles: + - merge.coalesce: + input: + - Col1 + - Col2 + output: Result + """, + dataframe=pd.DataFrame({ + 'Col1': [0, 1, None], + 'Col2': [99, 99, 99 ], + }) + ) + assert df['Result'].tolist() == [0, 1, 99] + + def test_false_not_skipped(self): + """ + Test that boolean False in a column is returned instead of falling back to the next column + """ + df = wrangles.recipe.run( + """ + wrangles: + - merge.coalesce: + input: + - Col1 + - Col2 + output: Result + """, + dataframe=pd.DataFrame({ + 'Col1': [False, True, None ], + 'Col2': ['fallback', 'fallback', 'fallback'], + }) + ) + assert df['Result'].tolist() == [False, True, 'fallback'] + + def test_lists_falsy_values_not_skipped(self): + """ + Test that falsy values (0, 0.0, False) inside a per-cell list are returned + as the first valid element rather than being skipped + """ + df = wrangles.recipe.run( + """ + wrangles: + - merge.coalesce: + input: Col + """, + dataframe=pd.DataFrame({ + 'Col': [ + [0, 1, 2 ], + [0.0, 1.5, 2.5], + [False, True ], + ['', 'b', 'c' ], + ] + }) + ) + assert df['Col'].tolist() == [0, 0.0, False, 'b'] + + class TestMergeConcatenate: """ All concatenate tests diff --git a/tests/recipes/wrangles/test_split.py b/tests/recipes/wrangles/test_split.py index c316f1952..509f9df87 100644 --- a/tests/recipes/wrangles/test_split.py +++ b/tests/recipes/wrangles/test_split.py @@ -1281,6 +1281,175 @@ def test_split_dictionary_empty(self): ) assert df.empty and df.columns.to_list() == ['Col'] + def test_split_dictionary_to_lists(self): + """ + Test splitting dictionary keys and values to parallel lists + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: + - Keys + - Values + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': [ + {"a": 123, "b": "kshdf"}, + {"b": 123, "c": "kshdf"} + ] + }) + ) + assert df['Keys'].to_list() == [["a", "b"], ["b", "c"]] + assert df['Values'].to_list() == [[123, "kshdf"], [123, "kshdf"]] + + def test_split_dictionary_to_lists_json(self): + """ + Test splitting JSON dictionary keys and values to lists + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: + - Keys + - Values + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': ['{"a": 123, "b": "kshdf"}'] + }) + ) + assert df['Keys'][0] == ["a", "b"] + assert df['Values'][0] == [123, "kshdf"] + + def test_split_dictionary_to_lists_multiple_inputs(self): + """ + Test splitting multiple dictionaries to keys and values lists + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: + - Dict1 + - Dict2 + output: + - Keys + - Values + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'Dict1': [{"a": 1, "b": 2}], + 'Dict2': [{"b": 3, "c": 4}] + }) + ) + assert df['Keys'][0] == ["a", "b", "c"] + assert df['Values'][0] == [1, 3, 4] + + def test_split_dictionary_to_lists_where(self): + """ + Test split.dictionary output_format to_lists using where + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: + - Keys + - Values + output_format: to_lists + where: numbers > 3 + """, + dataframe=pd.DataFrame({ + 'My Dict': [ + {"a": 1}, + {"b": 2}, + {"c": 3} + ], + 'numbers': [3, 4, 5] + }) + ) + assert df['Keys'].to_list() == ["", ["b"], ["c"]] + assert df['Values'].to_list() == ["", [2], [3]] + + def test_split_dictionary_to_lists_default_output(self): + """ + Test split.dictionary output_format to_lists default output columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': [{"a": 1}] + }) + ) + assert df['Keys'][0] == ["a"] + assert df['Values'][0] == [1] + + def test_split_dictionary_to_lists_output_error(self): + """ + Test split.dictionary output_format to_lists validates output column count + """ + with pytest.raises(ValueError, match="exactly two output columns"): + wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: Keys + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': [{"a": 1}] + }) + ) + + def test_split_dictionary_invalid_output_format(self): + """ + Test split.dictionary validates output_format + """ + with pytest.raises(ValueError, match="output_format must be one of"): + wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output_format: invalid + """, + dataframe=pd.DataFrame({ + 'My Dict': [{"a": 1}] + }) + ) + + def test_split_dictionary_to_lists_empty(self): + """ + Test split.dictionary output_format to_lists with an empty column + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: + - Keys + - Values + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': [] + }) + ) + assert df.empty and df.columns.to_list() == ['My Dict', 'Keys', 'Values'] + class TestTokenize: """ diff --git a/tests/test_wrangles.py b/tests/test_wrangles.py index fdc71f4c3..759319189 100644 --- a/tests/test_wrangles.py +++ b/tests/test_wrangles.py @@ -119,7 +119,6 @@ def test_extract_error_3(): def test_extract_html_str(): result = wrangles.extract.html('Wrangle Works!', dataType='text') assert result == 'Wrangle Works!' -# Translate def test_translate(mocker): translate_module = importlib.import_module('wrangles.translate') mocker.patch.object( @@ -133,6 +132,7 @@ def test_translate(mocker): assert result != 'My name is Chris' assert 'Chris' in result + def test_translate_list(mocker): translate_module = importlib.import_module('wrangles.translate') mocker.patch.object( @@ -141,10 +141,10 @@ def test_translate_list(mocker): return_value=['Ich heisse Chris'], ) result = wrangles.translate(['My name is Chris'], 'DE') - assert isinstance(result, list) - assert result[0] - assert result[0] != 'My name is Chris' + assert isinstance(result, list) and len(result) == 1 + assert isinstance(result[0], str) assert 'Chris' in result[0] + assert result[0] != 'My name is Chris' # Invalid input type (dict) def test_translate_typeError(): @@ -154,15 +154,15 @@ def test_translate_typeError(): def test_translate_list_lower_case(): result = wrangles.translate(['PRUEBA UNO'], 'EN-GB', case= 'lower') - assert result[0] == 'test one' - + assert 'test' in result[0].lower() and 'one' in result[0].lower() + def test_translate_list_upper_case(): result = wrangles.translate(['prueba dos'], 'EN-GB', case= 'upper') - assert result[0] == 'TEST TWO' + assert 'test' in result[0].lower() and 'two' in result[0].lower() def test_translate_list_title_case(): result = wrangles.translate(['PRueBa TrEs'], 'EN-GB', case= 'title') - assert result[0] == 'Test Three' + assert 'test' in result[0].lower() and 'three' in result[0].lower() # Standardize @@ -307,6 +307,35 @@ def test_lookup_list_value_list_column(): result = wrangles.lookup(["a"], "fe730444-1bda-4fcd", ["Value"]) assert result == [[1]] +def test_lookup_n_single_input(): + """ + Test n returns a list of n matches for a single input + """ + result = wrangles.lookup("Rachel", "e8658a6f-c694-45d0", n=2) + assert isinstance(result, list) + assert len(result) == 2 + +def test_lookup_n_list_input(): + """ + Test n returns a list of n-lists for a list of inputs + """ + result = wrangles.lookup(["Rachel", "Dolores"], "e8658a6f-c694-45d0", n=2) + assert isinstance(result, list) + assert len(result) == 2 + assert isinstance(result[0], list) and len(result[0]) == 2 + assert isinstance(result[1], list) and len(result[1]) == 2 + +def test_lookup_n_with_column(): + """ + Test n with a specific column still returns a list of n dicts - + requesting a single column does not collapse the dict to that + column's value when n > 1 + """ + result = wrangles.lookup("Rachel", "e8658a6f-c694-45d0", "Value", n=2) + assert isinstance(result, list) + assert len(result) == 2 + assert all(isinstance(match, dict) and "Value" in match for match in result) + def test_embedding_single(): """ Test generating an embedding from a single value diff --git a/wrangles/batching.py b/wrangles/batching.py index 3e9e64cb4..da2f24b9f 100644 --- a/wrangles/batching.py +++ b/wrangles/batching.py @@ -3,6 +3,7 @@ are unable to be processed in a single request """ +import logging as _logging from . import auth as _auth from . import utils as _utils @@ -16,8 +17,13 @@ def batch_api_calls(url, params, input_list, batch_size): # immediately return an empty list return [] + total_batches = (len(input_list) + batch_size - 1) // batch_size + _logging.info(f": Starting batch API calls :: url :: {url}, batch_size :: {batch_size}, total_records :: {len(input_list)}") + results = None for i in range(0, len(input_list), batch_size): + batch_num = i // batch_size + 1 + _logging.debug(f": Processing batch {batch_num} of {total_batches}") headers = {'Authorization': f'Bearer {_auth.get_access_token()}'} response = _utils.request_retries( request_type='POST', @@ -31,6 +37,7 @@ def batch_api_calls(url, params, input_list, batch_size): # Checking status code if str(response.status_code)[0] != '2': + _logging.error(f": Batch API request failed :: status_code :: {response.status_code}, batch :: {batch_num}") raise ValueError(f"Status Code: {response.status_code} - {response.reason}. {response.text} \n") response_json = response.json() diff --git a/wrangles/classify.py b/wrangles/classify.py index f1719e16b..ea17d3ebb 100644 --- a/wrangles/classify.py +++ b/wrangles/classify.py @@ -2,6 +2,7 @@ Functions to classify information. """ from typing import Union as _Union +import logging as _logging from . import config as _config from . import data as _data from . import batching as _batching @@ -51,6 +52,7 @@ def classify( if purpose != 'classify': raise ValueError(f'Using {purpose} model_id {model_id} in a classify function.') + _logging.info(f": Classifying input :: model_id :: {model_id}, record_count :: {len(json_data)}") results = _batching.batch_api_calls(url, params, json_data, batch_size) if isinstance(input, str): diff --git a/wrangles/compare.py b/wrangles/compare.py index 740a44891..38a840263 100644 --- a/wrangles/compare.py +++ b/wrangles/compare.py @@ -4,6 +4,7 @@ from collections import OrderedDict as _OrderedDict from difflib import SequenceMatcher as _SequenceMatcher +import logging as _logging import unicodedata from typing import Tuple @@ -138,6 +139,7 @@ def contrast(input: list, type: str ='difference', char: str = ' ', case_sensiti :param char: The character to split the strings on. Default is a space :param case_sensitive: Whether the comparison is case sensitive. Default is True """ + _logging.debug(f": Comparing {len(input)} records :: type :: {type}, case_sensitive :: {case_sensitive}") results = [] for row in input: @@ -195,6 +197,7 @@ def overlap( Find the matching characters between two strings. return: 2D list with the matched elements or the matched elements and the ratio of similarity in a list """ + _logging.debug(f": Computing overlap for {len(input)} record pairs") results = [] for row in input: @@ -269,6 +272,7 @@ def overlap( return results def deduplicate(result, enabled=False, ignore_case=False): + _logging.debug(f": Deduplicating {len(result)} items :: ignore_case :: {ignore_case}") if not enabled: return result diff --git a/wrangles/connectors/__init__.py b/wrangles/connectors/__init__.py index e8039f7f4..1c50c4c0e 100644 --- a/wrangles/connectors/__init__.py +++ b/wrangles/connectors/__init__.py @@ -3,8 +3,10 @@ """ from . import akeneo +from . import access from . import ckan from . import concurrent +from . import duckdb from . import excel from . import file from . import http @@ -26,4 +28,4 @@ from . import train from . import jinja from . import _formatting -from . import input \ No newline at end of file +from . import input diff --git a/wrangles/connectors/access.py b/wrangles/connectors/access.py new file mode 100644 index 000000000..f7485f445 --- /dev/null +++ b/wrangles/connectors/access.py @@ -0,0 +1,298 @@ +""" +Connector to read/write from Microsoft Access databases using ODBC. +""" +import logging as _logging +from typing import Union as _Union + +import pandas as _pd +from pandas.api import types as _pd_types + +from ..utils import ( + LazyLoader as _LazyLoader, + wildcard_expansion as _wildcard_expansion, +) + + +_pyodbc = _LazyLoader('pyodbc') + +_schema = {} + + +def _quote_identifier(identifier: str) -> str: + return f"[{str(identifier).replace(']', ']]')}]" + + +def _connection_string( + database: str = None, + connection_string: str = None, + driver: str = 'Microsoft Access Driver (*.mdb, *.accdb)', + password: str = None, +) -> str: + if connection_string: + return connection_string + if database is None: + raise ValueError('database or connection_string must be provided') + + conn = f"DRIVER={{{driver}}};DBQ={database};" + if password: + conn += f"PWD={password};" + return conn + + +def _access_type(dtype) -> str: + if _pd_types.is_bool_dtype(dtype): + return 'BIT' + if _pd_types.is_integer_dtype(dtype): + return 'INTEGER' + if _pd_types.is_float_dtype(dtype): + return 'DOUBLE' + if _pd_types.is_datetime64_any_dtype(dtype): + return 'DATETIME' + return 'LONGTEXT' + + +def _table_exists(cursor, table: str) -> bool: + return cursor.tables(table=table, tableType='TABLE').fetchone() is not None + + +def _create_table(cursor, table: str, df: _pd.DataFrame) -> None: + columns = ', '.join( + f"{_quote_identifier(column)} {_access_type(dtype)}" + for column, dtype in df.dtypes.items() + ) + cursor.execute(f"CREATE TABLE {_quote_identifier(table)} ({columns})") + + +def read( + command: str, + database: str = None, + connection_string: str = None, + driver: str = 'Microsoft Access Driver (*.mdb, *.accdb)', + password: str = None, + columns: _Union[str, list] = None, + params: _Union[list, tuple] = None, + **kwargs +) -> _pd.DataFrame: + """ + Read data from a Microsoft Access database. + + >>> from wrangles.connectors import access + >>> df = access.read(database='database.accdb', command='SELECT * FROM table') + + :param command: SQL command to select data. + :param database: Access database file path. Not required if connection_string is supplied. + :param connection_string: Full ODBC connection string. If provided, database, driver, and password are ignored. + :param driver: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + :param password: Optional database password. + :param columns: (Optional) Subset of columns to be returned. This is less efficient than specifying in the SQL command. + :param params: (Optional) Variables to pass to a parameterized query. + """ + target = database or connection_string + _logging.info(f": Reading data from Microsoft Access :: {target}") + + conn_string = _connection_string(database, connection_string, driver, password) + with _pyodbc.connect(conn_string) as conn: + df = _pd.read_sql(command, conn, params=params, **kwargs) + + if columns is not None: + columns = _wildcard_expansion(df.columns, columns) + df = df[columns] + + return df + + +_schema['read'] = r""" +type: object +description: Import data from a Microsoft Access Database +required: + - command +properties: + database: + type: string + description: Access database file path. Not required if connection_string is supplied. + connection_string: + type: string + description: Full ODBC connection string. If provided, database, driver, and password are ignored. + driver: + type: string + description: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + password: + type: string + description: Optional database password. + command: + type: string + description: |- + SQL command to select data. + Note - using variables here can make your recipe vulnerable + to sql injection. Use params if using variables from + untrusted sources. + columns: + type: + - string + - array + description: A list with a subset of the columns to import. This is less efficient than specifying in the command. + params: + type: array + description: Variables to pass to a parameterized query. +""" + + +def write( + df: _pd.DataFrame, + table: str, + database: str = None, + connection_string: str = None, + driver: str = 'Microsoft Access Driver (*.mdb, *.accdb)', + password: str = None, + action: str = 'INSERT', + columns: _Union[str, list] = None, +) -> None: + """ + Write data to a Microsoft Access database. + + >>> from wrangles.connectors import access + >>> access.write(df, database='database.accdb', table='table') + + :param df: Pandas Dataframe to be written. + :param table: Table to be exported to. + :param database: Access database file path. Not required if connection_string is supplied. + :param connection_string: Full ODBC connection string. If provided, database, driver, and password are ignored. + :param driver: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + :param password: Optional database password. + :param action: INSERT appends, REPLACE recreates the table, FAIL errors if the table exists. Defaults to INSERT. + :param columns: (Optional) Subset of the columns to be written. If not provided, all columns will be output. + """ + target = database or connection_string + _logging.info(f": Writing data to Microsoft Access :: {target} / {table}") + + action = action.upper() + if action not in ('INSERT', 'REPLACE', 'FAIL'): + raise ValueError('Invalid action. Expected INSERT, REPLACE, or FAIL.') + + if columns is not None: + columns = _wildcard_expansion(df.columns, columns) + df = df[columns] + + conn_string = _connection_string(database, connection_string, driver, password) + with _pyodbc.connect(conn_string) as conn: + cursor = conn.cursor() + exists = _table_exists(cursor, table) + + if action == 'FAIL' and exists: + raise ValueError(f"Table already exists: {table}") + if action == 'REPLACE' and exists: + cursor.execute(f"DROP TABLE {_quote_identifier(table)}") + exists = False + if not exists: + _create_table(cursor, table, df) + + if not df.empty: + table_name = _quote_identifier(table) + column_names = ', '.join(_quote_identifier(column) for column in df.columns) + placeholders = ', '.join('?' for _ in df.columns) + sql = f"INSERT INTO {table_name} ({column_names}) VALUES ({placeholders})" + cursor.executemany(sql, df.where(_pd.notnull(df), None).itertuples(index=False, name=None)) + + conn.commit() + + +_schema['write'] = """ +type: object +description: Export data to a Microsoft Access Database +required: + - table +properties: + database: + type: string + description: Access database file path. Not required if connection_string is supplied. + connection_string: + type: string + description: Full ODBC connection string. If provided, database, driver, and password are ignored. + driver: + type: string + description: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + password: + type: string + description: Optional database password. + table: + type: string + description: The table to write to + action: + type: string + description: INSERT appends, REPLACE recreates the table, FAIL errors if the table exists. Defaults to INSERT. + enum: + - INSERT + - REPLACE + - FAIL + columns: + type: + - string + - array + description: A list of the columns to write to the table. If omitted, all columns will be written. +""" + + +def run( + command: _Union[str, list], + database: str = None, + connection_string: str = None, + driver: str = 'Microsoft Access Driver (*.mdb, *.accdb)', + password: str = None, + params: _Union[list, tuple] = None, +) -> None: + """ + Run a command on a Microsoft Access database. + + >>> wrangles.connectors.access.run( + >>> database='database.accdb', + >>> command='' + >>> ) + + :param command: SQL command or a list of SQL commands to execute. + :param database: Access database file path. Not required if connection_string is supplied. + :param connection_string: Full ODBC connection string. If provided, database, driver, and password are ignored. + :param driver: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + :param password: Optional database password. + :param params: Variables to pass to a parameterized query. + """ + target = database or connection_string + _logging.info(f": Executing Microsoft Access Command :: {target}") + + if isinstance(command, str): + command = [command] + + conn_string = _connection_string(database, connection_string, driver, password) + with _pyodbc.connect(conn_string) as conn: + cursor = conn.cursor() + for sql in command: + cursor.execute(sql, params or ()) + conn.commit() + + +_schema['run'] = r""" +type: object +description: Run a command against a Microsoft Access Database +required: + - command +properties: + database: + type: string + description: Access database file path. Not required if connection_string is supplied. + connection_string: + type: string + description: Full ODBC connection string. If provided, database, driver, and password are ignored. + driver: + type: string + description: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + password: + type: string + description: Optional database password. + command: + type: + - string + - array + description: SQL command or a list of SQL commands to execute + params: + type: array + description: Variables to pass to a parameterized query. +""" diff --git a/wrangles/connectors/akeneo.py b/wrangles/connectors/akeneo.py index df7ceda59..65b1b0d9c 100644 --- a/wrangles/connectors/akeneo.py +++ b/wrangles/connectors/akeneo.py @@ -49,13 +49,9 @@ def read( :return: A Pandas dataframe of the returned results """ _logging.info(f": Reading data from Akeneo :: {host} / {source}") - if parameters is None: - parameters = {} - # Set to max temporarily - parameters['limit'] = 100 + parameters = {**(parameters or {}), 'limit': 100} - # TODO: deal with errors appropriately - token = _requests.post( + auth_response = _requests.post( f"{host}/api/oauth/v1/token", auth = (client_id, client_secret), json = { @@ -63,18 +59,34 @@ def read( "password" : password, "grant_type" : "password" } - ).json()['access_token'] - - # TODO: Needs to deal with pagination - data = _requests.get( - f"{host}/api/rest/v1/{source}", - params=parameters, - headers={ - 'Accept': 'application/json', - 'Authorization': f"Bearer {token}" - } - ).json()['_embedded']['items'] - + ) + if not auth_response.ok: + try: + message = auth_response.json().get('message', auth_response.text) + except Exception: + message = auth_response.text + raise ValueError(f"Akeneo authentication failed: {message}") + + token = auth_response.json()['access_token'] + headers={ + 'Accept': 'application/json', + 'Authorization': f"Bearer {token}" + } + + data = [] + url = f"{host}/api/rest/v1/{source}" + params = parameters + + while url: + response = _requests.get(url, params=params, headers=headers) + if not response.ok: + json_response = response.json() + raise ValueError(f"Status Code: {json_response.get('code', response.status_code)} Message: {json_response.get('message', response.text)}") + response_json = response.json() + data.extend(response_json['_embedded']['items']) + url = response_json.get('_links', {}).get('next', {}).get('href') + params = None + df = _pd.json_normalize(data, max_level=0) if columns: @@ -180,16 +192,23 @@ def write( """ _logging.info(f": Writing data to Akeneo :: {host} / {source}") - # TODO: handle errors appropriately - token = _requests.post( + auth_response = _requests.post( f"{host}/api/oauth/v1/token", - auth = (client_id, client_secret), + auth=(client_id, client_secret), json={ - "username" : user, - "password" : password, - "grant_type" : "password" + "username": user, + "password": password, + "grant_type": "password" } - ).json()['access_token'] + ) + if not auth_response.ok: + try: + message = auth_response.json().get('message', auth_response.text) + except Exception: + message = auth_response.text + raise ValueError(f"Akeneo authentication failed: {message}") + + token = auth_response.json()['access_token'] # TODO: batch this if required?? # Create payload for Akeneo diff --git a/wrangles/connectors/duckdb.py b/wrangles/connectors/duckdb.py new file mode 100644 index 000000000..22d1115c4 --- /dev/null +++ b/wrangles/connectors/duckdb.py @@ -0,0 +1,205 @@ +""" +Connector to read/write from DuckDB database files. +""" +import logging as _logging +from typing import Union as _Union + +import pandas as _pd + +from ..utils import ( + LazyLoader as _LazyLoader, + wildcard_expansion as _wildcard_expansion, +) + + +_duckdb = _LazyLoader('duckdb') + +_schema = {} + + +def _quote_identifier(identifier: str) -> str: + return '.'.join( + f'"{part.replace(chr(34), chr(34) * 2)}"' + for part in str(identifier).split('.') + ) + + +def read( + database: str, + command: str, + columns: _Union[str, list] = None, + params: _Union[list, tuple, dict] = None, + **kwargs +) -> _pd.DataFrame: + """ + Read data from a DuckDB database. + + >>> from wrangles.connectors import duckdb + >>> df = duckdb.read(database='database.duckdb', command='SELECT * FROM table') + + :param database: The database to connect to including the file path. Use ':memory:' for an in-memory database. + :param command: SQL command to select data. + :param columns: (Optional) Subset of columns to be returned. This is less efficient than specifying in the SQL command. + :param params: (Optional) Variables to pass to a parameterized query. + """ + _logging.info(f": Reading data from DuckDB :: {database}") + + with _duckdb.connect(database=database, **kwargs) as conn: + df = conn.execute(command, params or ()).fetchdf() + + if columns is not None: + columns = _wildcard_expansion(df.columns, columns) + df = df[columns] + + return df + + +_schema['read'] = r""" +type: object +description: Import data from a DuckDB Database +required: + - database + - command +properties: + database: + type: string + description: The database to connect to including the file path. Use ':memory:' for an in-memory database. + command: + type: string + description: |- + SQL command to select data. + Note - using variables here can make your recipe vulnerable + to sql injection. Use params if using variables from + untrusted sources. + columns: + type: + - string + - array + description: A list with a subset of the columns to import. This is less efficient than specifying in the command. + params: + type: + - array + - object + description: Variables to pass to a parameterized query. +""" + + +def write( + df: _pd.DataFrame, + database: str, + table: str, + action: str = 'INSERT', + columns: _Union[str, list] = None, + **kwargs +) -> None: + """ + Write data to a DuckDB database. + + >>> from wrangles.connectors import duckdb + >>> duckdb.write(df, database='database.duckdb', table='table') + + :param df: Pandas Dataframe to be written. + :param database: The database to connect to including the file path. Use ':memory:' for an in-memory database. + :param table: Table to be exported to. + :param action: INSERT appends, REPLACE recreates the table, FAIL errors if the table exists. Defaults to INSERT. + :param columns: (Optional) Subset of the columns to be written. If not provided, all columns will be output. + """ + _logging.info(f": Writing data to DuckDB :: {database} / {table}") + + action = action.upper() + if columns is not None: + columns = _wildcard_expansion(df.columns, columns) + df = df[columns] + + table_name = _quote_identifier(table) + with _duckdb.connect(database=database, **kwargs) as conn: + conn.register('_wrangles_df', df) + + if action == 'FAIL': + conn.execute(f'CREATE TABLE {table_name} AS SELECT * FROM _wrangles_df') + elif action == 'REPLACE': + conn.execute(f'CREATE OR REPLACE TABLE {table_name} AS SELECT * FROM _wrangles_df') + elif action == 'INSERT': + conn.execute(f'CREATE TABLE IF NOT EXISTS {table_name} AS SELECT * FROM _wrangles_df WHERE false') + conn.execute(f'INSERT INTO {table_name} SELECT * FROM _wrangles_df') + else: + raise ValueError('Invalid action. Expected INSERT, REPLACE, or FAIL.') + + +_schema['write'] = """ +type: object +description: Export data to a DuckDB Database +required: + - database + - table +properties: + database: + type: string + description: The database to connect to including the file path. Use ':memory:' for an in-memory database. + table: + type: string + description: The table to write to + action: + type: string + description: INSERT appends, REPLACE recreates the table, FAIL errors if the table exists. Defaults to INSERT. + enum: + - INSERT + - REPLACE + - FAIL + columns: + type: + - string + - array + description: A list of the columns to write to the table. If omitted, all columns will be written. +""" + + +def run( + database: str, + command: _Union[str, list], + params: _Union[list, tuple, dict] = None, + **kwargs +) -> None: + """ + Run a command on a DuckDB database. + + >>> wrangles.connectors.duckdb.run( + >>> database='database.duckdb', + >>> command='' + >>> ) + + :param database: The database to connect to including the file path. Use ':memory:' for an in-memory database. + :param command: SQL command or a list of SQL commands to execute. + :param params: Variables to pass to a parameterized query. + """ + _logging.info(f": Executing DuckDB Command :: {database}") + + if isinstance(command, str): + command = [command] + + with _duckdb.connect(database=database, **kwargs) as conn: + for sql in command: + conn.execute(sql, params or ()) + + +_schema['run'] = r""" +type: object +description: Run a command against a DuckDB Database +required: + - database + - command +properties: + database: + type: string + description: The database to connect to including the file path. Use ':memory:' for an in-memory database. + command: + type: + - string + - array + description: SQL command or a list of SQL commands to execute + params: + type: + - array + - object + description: Variables to pass to a parameterized query. +""" diff --git a/wrangles/connectors/file.py b/wrangles/connectors/file.py index 0781a14a4..370e0ac86 100644 --- a/wrangles/connectors/file.py +++ b/wrangles/connectors/file.py @@ -20,6 +20,25 @@ _schema = {} +_ILLEGAL_CHARACTERS_RE = _re.compile( + '[^\u0009\u000A\u000D\u0020-\uD7FF\uE000-\uFFFD\U00010000-\U0010FFFF]' +) + + +def _clean_cell_value(value): + if isinstance(value, str): + return _ILLEGAL_CHARACTERS_RE.sub('', value) + elif isinstance(value, list): + return [_clean_cell_value(v) for v in value] + elif isinstance(value, dict): + return {k: _clean_cell_value(v) for k, v in value.items()} + return value + + +def _remove_illegal_characters(df: _pd.DataFrame) -> _pd.DataFrame: + return df.map(_clean_cell_value) + + def read( name: str, columns: _Union[str, list] = None, @@ -46,6 +65,10 @@ def read( :param kwargs: (Optional) Named arguments to pass to respective pandas function. :return: A Pandas dataframe of the imported data. """ + # Accept path-like objects (e.g. pathlib.Path) by converting to str + if isinstance(name, _os.PathLike): + name = str(name) + if ( isinstance(name, dict) and all(x in name for x in ['name', 'data', 'mimeType']) @@ -114,7 +137,7 @@ def read( properties: name: type: string - description: The name of the file to import + description: The name or path of the file to import. Accepts a string or a Path object (pathlib.Path / os.PathLike). columns: type: array description: Columns to select @@ -187,6 +210,10 @@ def write(df: _pd.DataFrame, name: str, columns: _Union[str, list] = None, file_ :param formatting: (Optional) A dictionary of formatting options to apply to Excel files. :param kwargs: (Optional) Named arguments to pass to respective pandas function. """ + # Accept path-like objects (e.g. pathlib.Path) by converting to str + if isinstance(name, _os.PathLike): + name = str(name) + _logging.info(f": Writing data to file :: {name}") # Select only specific columns if user requests them @@ -223,7 +250,7 @@ def write(df: _pd.DataFrame, name: str, columns: _Union[str, list] = None, file_ else: with _pd.ExcelWriter(file_object, engine='openpyxl') as writer: - df.to_excel(writer, **kwargs) + _remove_illegal_characters(df).to_excel(writer, **kwargs) worksheet = writer.sheets[kwargs.get('sheet_name', 'Sheet1')] @@ -282,7 +309,7 @@ def write(df: _pd.DataFrame, name: str, columns: _Union[str, list] = None, file_ properties: name: type: string - description: The name of the file to write. + description: The name or path of the file to write. Accepts a string or a Path object (pathlib.Path / os.PathLike). columns: type: array description: A list of the columns to write. If omitted, all columns will be written. diff --git a/wrangles/connectors/http.py b/wrangles/connectors/http.py index 2822019c1..2d0554692 100644 --- a/wrangles/connectors/http.py +++ b/wrangles/connectors/http.py @@ -15,6 +15,7 @@ def _get_oauth_token(url, method="POST", **kwargs): :param method: The http method to use. Default POST. :return: The OAuth token """ + _logging.debug(f": Fetching OAuth token :: url :: {url}, method :: {method}") response = _requests.request(url=url, method=method, **kwargs) if not response.ok: raise RuntimeError( @@ -62,6 +63,7 @@ def run( **kwargs ) if not response.ok: + _logging.error(f": HTTP request failed :: status :: {response.status_code}, url :: {url}") raise RuntimeError( f"Request failed with status code {response.status_code}. Response: {response.text}" ) diff --git a/wrangles/connectors/jinja.py b/wrangles/connectors/jinja.py index 083149598..d283284e8 100644 --- a/wrangles/connectors/jinja.py +++ b/wrangles/connectors/jinja.py @@ -1,6 +1,7 @@ """ Use JINJA to manipulate a template """ +import logging as _logging from ..utils import LazyLoader as _LazyLoader # Lazy load external dependency @@ -17,6 +18,7 @@ def run(template: dict, context: dict, output_file: str): :param context: A dictionary used to define the output template :param output_file: File name/path for the file to be output """ + _logging.info(f": Processing Jinja template :: output_file :: {output_file}") if 'file' in template: template = _jinja.Environment( loader=_jinja.FileSystemLoader(''), diff --git a/wrangles/connectors/memory.py b/wrangles/connectors/memory.py index 0c5b5a17d..e74b5b31b 100644 --- a/wrangles/connectors/memory.py +++ b/wrangles/connectors/memory.py @@ -25,6 +25,7 @@ def clear(): Clear and reset any existing data stored in the connector """ + _logging.info(": Clearing memory connector") global dataframes, variables, queue dataframes = {} variables = {} diff --git a/wrangles/connectors/mysql.py b/wrangles/connectors/mysql.py index 703c828fe..1980e6ec7 100644 --- a/wrangles/connectors/mysql.py +++ b/wrangles/connectors/mysql.py @@ -4,8 +4,14 @@ import pandas as _pd from typing import Union as _Union import logging as _logging -from ..utils import wildcard_expansion as _wildcard_expansion - +from ..utils import ( + wildcard_expansion as _wildcard_expansion, + LazyLoader as _LazyLoader +) + +# Lazy load optional dependencies +_pymysql = _LazyLoader('pymysql') +_sqlalchemy = _LazyLoader('sqlalchemy') _schema = {} diff --git a/wrangles/connectors/notification.py b/wrangles/connectors/notification.py index d45eb2533..f82d8ce98 100644 --- a/wrangles/connectors/notification.py +++ b/wrangles/connectors/notification.py @@ -2,6 +2,7 @@ Send notifications to a varity of services """ from typing import Union as _Union +import logging as _logging from ..utils import LazyLoader as _LazyLoader # Lazy load external dependency @@ -24,6 +25,7 @@ def run( :param body: The body of the notification :param attachment: A file path & name to attach to the message. Supports a single file or a list of files. Must be supported by the specific notification type. """ + _logging.info(f": Sending notification :: title :: {title}") app_object = _apprise.Apprise() app_object.add(url) app_object.notify( @@ -121,6 +123,7 @@ def run( :param attachment: A file path & name to attach to the message. Supports a single file or a list of files. Must be supported by the specific notification type. :param format: The format of the message. One of 'text', 'markdown' or 'html'. Default is 'text' """ + _logging.info(f": Sending Telegram message :: chat_id :: {chat_id}") url = f"tgram://{bot_token}/{chat_id}/?format={format}" run(url, title, body, attachment) diff --git a/wrangles/connectors/s3.py b/wrangles/connectors/s3.py index 2b3c3cf16..51b6b5e57 100644 --- a/wrangles/connectors/s3.py +++ b/wrangles/connectors/s3.py @@ -2,6 +2,7 @@ from io import BytesIO as _BytesIO from typing import Union as _Union import pandas as _pd +import pathlib as _pathlib from . import file as _file from ..utils import LazyLoader as _LazyLoader import os as _os @@ -267,16 +268,16 @@ class upload_files: bucket: type: string description: S3 Bucket - file_key: + save_as: type: - string - array - description: S3 file key or list of keys to upload as. - save_as: + description: S3 file key(s) to upload as. Include directory in this path. + file: type: - string - array - description: File or list of files to upload (replaces 'file'). + description: File or list of files to upload. Accepts strings or pathlib.Path objects. endpoint_url: type: string description: Override the S3 host for alternative S3 storage providers. @@ -289,46 +290,52 @@ class upload_files: """ } - def run(bucket: str, save_as: _Union[str, list] = None, file_key: _Union[str, list] = None, **kwargs): + def run(bucket: str, file: _Union[str, list] = None, save_as: _Union[str, list] = None, **kwargs): """ Upload file(s) to S3 from the local file system. :param bucket: S3 Bucket - :param save_as: File or list of files to upload. - :param file_key: S3 file key or list of keys to upload as. + :param file: File or list of files to upload. Accepts strings or pathlib.Path objects. + :param save_as: S3 file key(s) to upload as. Include directory in this path. :param endpoint_url: Override the S3 host for alternative S3 storage providers. :param aws_access_key_id: Set the access key. Can also be set as an environment variable :param aws_secret_access_key: Set the access secret. Can also be set as an environment variable """ # Backwards compatibility: accept deprecated 'key' via kwargs compat_key = kwargs.pop('key', None) - if file_key is None and compat_key is not None: - file_key = compat_key + if save_as is None and compat_key is not None: + save_as = compat_key # Backwards compatibility: accept deprecated 'file' via kwargs compat_file = kwargs.pop('file', None) - if save_as is None and compat_file is not None: - save_as = compat_file - if save_as is None: - raise ValueError("save_as must be provided") - - _logging.info(f": Uploading files to S3 :: {bucket} / {file_key}") + if file is None and compat_file is not None: + file = compat_file + if file is None: + raise ValueError("file must be provided") + + # Normalize Path objects to strings + if isinstance(file, _pathlib.Path): + file = str(file) + elif isinstance(file, list): + file = [str(f) if isinstance(f, _pathlib.Path) else f for f in file] + + _logging.info(f": Uploading files to S3 :: {bucket} / {save_as}") s3 = _boto3.client('s3', **kwargs) - if isinstance(save_as, str): save_as = [save_as] + if isinstance(file, str): file = [file] # If a list of filename isn't provided, then save # in the current directory as the file_key's filename - if not file_key: - file_key = [k.split('/')[-1] for k in save_as] + if not save_as: + save_as = [k.split('/')[-1] for k in file] - if isinstance(file_key, str): - file_key = [file_key] + if isinstance(save_as, str): + save_as = [save_as] - if len(save_as) != len(file_key): + if len(file) != len(save_as): raise ValueError('s3.upload_files: An equal number of files and keys must be provided') - for f, k in zip(save_as, file_key): + for f, k in zip(file, save_as): try: s3.upload_file(f, bucket, k) except s3.exceptions.ClientError as e: diff --git a/wrangles/connectors/sftp.py b/wrangles/connectors/sftp.py index 2e6925d01..e0e7f3783 100644 --- a/wrangles/connectors/sftp.py +++ b/wrangles/connectors/sftp.py @@ -7,13 +7,14 @@ import pandas as _pd import io as _io import logging as _logging + from . import file as _file from ..utils import LazyLoader as _LazyLoader # Lazy load external dependencies _fabric = _LazyLoader('fabric') -_RSAKey = _LazyLoader('paramiko').RSAKey +_paramiko = _LazyLoader('paramiko') _schema = {} @@ -52,7 +53,7 @@ def read( if password: connect_kwargs = {'password': password, **connect_kwargs} elif pkey: - connect_kwargs = {'pkey': _RSAKey(file_obj=_io.StringIO(pkey)), **connect_kwargs} + connect_kwargs = {'pkey': _paramiko.RSAKey(file_obj=_io.StringIO(pkey)), **connect_kwargs} else: raise ValueError("Either password or pkey must be provided to connect to the SFTP server") @@ -162,7 +163,7 @@ def write( if password: connect_kwargs = {'password': password, **connect_kwargs} elif pkey: - connect_kwargs = {'pkey': _RSAKey(file_obj=_io.StringIO(pkey)), **connect_kwargs} + connect_kwargs = {'pkey': _paramiko.RSAKey(file_obj=_io.StringIO(pkey)), **connect_kwargs} else: raise ValueError("Either password or pkey must be provided to connect to the SFTP server") @@ -311,7 +312,7 @@ def run( if password: connect_kwargs = {'password': password, **connect_kwargs} elif pkey: - connect_kwargs = {'pkey': _RSAKey(file_obj=_io.StringIO(pkey)), **connect_kwargs} + connect_kwargs = {'pkey': _paramiko.RSAKey(file_obj=_io.StringIO(pkey)), **connect_kwargs} else: raise ValueError("Either password or pkey must be provided to connect to the SFTP server") @@ -418,7 +419,7 @@ def run( if password: connect_kwargs = {'password': password, **connect_kwargs} elif pkey: - connect_kwargs = {'pkey': _RSAKey(file_obj=_io.StringIO(pkey)), **connect_kwargs} + connect_kwargs = {'pkey': _paramiko.RSAKey(file_obj=_io.StringIO(pkey)), **connect_kwargs} else: raise ValueError("Either password or pkey must be provided to connect to the SFTP server") diff --git a/wrangles/extract.py b/wrangles/extract.py index 0f7a0d85a..5926f76b9 100644 --- a/wrangles/extract.py +++ b/wrangles/extract.py @@ -5,6 +5,7 @@ from typing import Union as _Union import concurrent.futures as _futures import json as _json +import logging as _logging import pandas as _pd from . import config as _config from . import data as _data @@ -33,6 +34,7 @@ def address( else: json_data = input + _logging.info(f": Extracting address {dataType} from {len(json_data)} records") url = f'{_config.api_host}/wrangles/extract/address' params = { 'responseFormat':'array', @@ -248,6 +250,25 @@ def _standardize_schema(node): for k, v in output.items() } + # Sanitize output keys: replace any character outside [a-zA-Z0-9_] with + # an underscore before sending to the model. Property names with special + # characters (e.g. parentheses, spaces) are sometimes altered by the model, + # producing key mismatches that cause silent empty columns. We remap the + # sanitized names back to the originals after receiving results. + _key_to_original = {} + _sanitized_output = {} + for _k, _v in output.items(): + _sk = _re.sub(r'[^a-zA-Z0-9_]', '_', _k) + _base, _n = _sk, 2 + while _sk in _key_to_original: + _sk = f"{_base}_{_n}" + _n += 1 + _key_to_original[_sk] = _k + _sanitized_output[_sk] = _v + _needs_remap = any(sk != ok for sk, ok in _key_to_original.items()) + if _needs_remap: + output = _sanitized_output + # Format any user submitted header messages if messages and not isinstance(messages, list): messages = [str(messages)] @@ -308,6 +329,7 @@ def _standardize_schema(node): **kwargs } + _logging.info(f": Extracting data using AI model :: model_id :: {model_id}, thread_count :: {threads}") with _futures.ThreadPoolExecutor(max_workers=threads) as executor: results = list(executor.map( _openai.chatGPT, @@ -319,6 +341,13 @@ def _standardize_schema(node): [retries] * len(input), )) + if _needs_remap: + results = [ + {_key_to_original.get(k, k): v for k, v in row.items()} + if isinstance(row, dict) else row + for row in results + ] + if input_was_scalar: if output_generic_key: return results[0].get('output', 'Failed') @@ -357,6 +386,7 @@ def attributes( else: json_data = input + _logging.info(f": Extracting attributes from {len(json_data)} records") url = f'{_config.api_host}/wrangles/extract/attributes' params = { 'responseFormat':'array', @@ -403,6 +433,7 @@ def codes( else: json_data = input + _logging.info(f": Extracting codes from {len(json_data)} records") url = f'{_config.api_host}/wrangles/extract/codes' params = {'responseFormat': 'array', **kwargs} batch_size = 10000 @@ -425,7 +456,9 @@ def custom( case_sensitive: bool = False, extract_raw: bool = False, use_spellcheck: bool = False, + include_empty_labels: bool = True, sort: str = 'training_order', + output_format: str = 'dict', **kwargs ) -> list: """ @@ -464,6 +497,15 @@ def custom( } model_properties = _data.model(model_id) + model_content = _data.model_content(model_id) + + model_labels = set() + for item in model_content['Data']: + if len(item) >= 2: + if ':' in item[1]: + label = item[1].split(':')[0] # Second column typically contains the label/type + model_labels.add(label.strip()) + # If model_id format is correct but no mode_id exists if model_properties.get('message', None) == 'error': raise ValueError('Incorrect model_id.\nmodel_id may be wrong or does not exists') @@ -509,13 +551,26 @@ def _entities_to_list(value): {results["columns"][i]: row[i] for i in range(len(row))} for row in results["data"] ] - if isinstance(results, list): if first_element and not use_labels: results = [x[0] if len(x) >= 1 else "" for x in results] - if use_labels and first_element: - results = [{k:v[0] for (k, v) in zip(objs.keys(), objs.values())} for objs in results] + if use_labels: + if include_empty_labels: + # Ensure every label has a key, create empty keys if missing. + # Use both labels discovered from results and labels defined in the model. + all_labels = set(model_labels or []) + for objs in results: + all_labels.update([str(k).lower() for k in objs.keys()]) + + for objs in results: + # Normalize existing keys to lower-case while preserving original keys + existing = {str(k).lower(): k for k in objs.keys()} + for label in all_labels: + if label not in existing: + objs[label] = [] + if first_element: + results = [{k: v[0] if isinstance(v, list) and v else "" for k, v in objs.items()} for objs in results] else: raise ValueError(f'API Response did not return an expected format for model {model_id}') @@ -543,6 +598,7 @@ def html( else: json_data = input + _logging.info(f": Extracting {dataType} from HTML") url = f'{_config.api_host}/wrangles/extract/html' params = { 'responseFormat': 'array', @@ -582,6 +638,7 @@ def properties( else: json_data = input + _logging.info(f": Extracting properties from {len(json_data)} records") url = f'{_config.api_host}/wrangles/extract/properties' params = {'responseFormat':'array', **kwargs} if type is not None: params['dataType'] = type @@ -619,6 +676,7 @@ def remove_words(input: _Union[str, list], to_remove: list, tokenize_to_remove: else: flags = 0 # this is the default for _re.sub + _logging.info(f": Removing words from {len(input)} records") results = [] for _in, _remove in zip(input, to_remove): @@ -669,7 +727,8 @@ def remove_words(input: _Union[str, list], to_remove: list, tokenize_to_remove: def brackets( input: str, find: list = _Union[str, list], - include_brackets: bool = False + include_brackets: bool = False, + return_data_type: str = "string" ) -> list: """ Extract values in brackets, [], {}, (), <> @@ -677,8 +736,9 @@ def brackets( :param input: Input string to search for brackets :param find: Types of brackets to find (e.g., 'round', 'square', 'curly', 'angled'). Default is all types. :param include_brackets: Whether to include brackets in the results - :return: List of lists of extracted values + :return: List of extracted values """ + _logging.info(": Extracting text from brackets") results = [] bracket_patterns = { 'round': r'\(.*?\)', @@ -704,6 +764,9 @@ def brackets( if include_brackets is False: re = [_re.sub(r'\[|\]|{|}|\(|\)|<|>', '', re[x]) for x in range(len(re))] - results.append(re) - + if return_data_type == "list": + results.append(re) + else: + results.append(', '.join(re)) + return results diff --git a/wrangles/format.py b/wrangles/format.py index 44550aeff..1f4097280 100644 --- a/wrangles/format.py +++ b/wrangles/format.py @@ -1,6 +1,7 @@ from typing import Union as _Union import types as _types import re as _re +import logging as _logging import pandas as _pandas import numpy as _np @@ -14,6 +15,7 @@ def concatenate(data_list, concat_char, skip_empty: bool=False): """ Concatenate a list of columns """ + _logging.debug(f": Concatenating {len(data_list)} records :: char :: {concat_char}") if skip_empty: return [ concat_char.join([str(x) for x in row if x]) @@ -46,6 +48,7 @@ def split( :param element: Slice the output lists to specific elements. :param skip_empty: If true, skip empty cells. """ + _logging.debug(f": Splitting {len(input_list)} records :: char :: {split_char}") # Split as either regex or simple string if split_char[:6] == 'regex:': split_char = split_char[6:].strip() @@ -122,15 +125,67 @@ def coalesce(input_list: list) -> list: Return the first not empty result for each row where each row has a list of possibilities """ + + _logging.debug(f": Coalescing {len(input_list)} values") + + if not input_list: + return [] + + # Use numpy optimization for rectangular (fixed-width) input (e.g. multi-column coalesce). + # Fall back to Python loop for ragged input (e.g. lists stored in a single column). + row_len = len(input_list[0]) + if row_len > 0 and all(len(row) == row_len for row in input_list): + arr = _np.empty((len(input_list), row_len), dtype=object) + for i, row in enumerate(input_list): + for j, val in enumerate(row): + arr[i, j] = val if val is not None else '' + + n_rows, m = arr.shape + + # Fast C-level check: find first column per row where value != '' + mask = arr != '' + idx = _np.argmax(mask, axis=1) + has_any = mask.any(axis=1) + result = arr[_np.arange(n_rows), idx].copy() + result[~has_any] = '' + + # arr != '' treats whitespace-only strings as non-empty; handle those edge cases + # by stripping only the selected values (1 per row) rather than the full matrix. + try: + stripped_result = _np.frompyfunc(str.strip, 1, 1)(result) + except (TypeError, AttributeError): + def _safe_strip(x): + try: + return x.strip() + except AttributeError: + return '' if not x else x + stripped_result = _np.frompyfunc(_safe_strip, 1, 1)(result) + needs_fallback = _np.where(has_any & (stripped_result == ''))[0] + for i in needs_fallback: + for j in range(int(idx[i]) + 1, m): + val = arr[i, j] + if isinstance(val, str): + val = val.strip() + if val: + result[i] = val + break + else: + result[i] = '' + + return result.tolist() + output_list = [] for row in input_list: output_row = '' for value in row: - if isinstance(value, str): value = value.strip() - if value: + if isinstance(value, str): + value = value.strip() + if value: + output_row = value + break + elif value is not None: output_row = value break - output_list.append(output_row) return output_list @@ -139,6 +194,7 @@ def price_breaks(df_input, header_cat, header_val): # pragma: no cover """ Rearrange price breaks """ + _logging.info(f": Processing price breaks for {len(df_input)} records") output = [] headers = [] i = 1 @@ -224,6 +280,7 @@ def remove_duplicates(input_list: list, ignore_case: bool = False) -> list: """ Remove duplicates from a list. Preserves input order. """ + _logging.debug(f": Removing duplicates :: ignore_case :: {ignore_case}") results = [] for row in input_list: # If row is a list, remove duplicates while ignoring case @@ -462,6 +519,7 @@ def tokenize( :param pattern: A custom regex pattern or regex string to split the input on :return: The tokenized list """ + _logging.debug(f": Tokenizing {len(input)} records") word_boundary_pattern = _re.compile(r"([\b\W\b])") def split_boundary_ignore_space(value): diff --git a/wrangles/generate.py b/wrangles/generate.py index ecfeeb9f3..f02ba7bfe 100644 --- a/wrangles/generate.py +++ b/wrangles/generate.py @@ -2,6 +2,7 @@ import concurrent.futures import copy import json +import logging as _logging from typing import Any, Dict, List, Literal, Union, Optional, Tuple import requests @@ -27,6 +28,7 @@ class PropertyDefinition(BaseModel): def _perform_web_search(query: str) -> str: + _logging.debug(f": Performing web search :: query :: {query}") if BeautifulSoup is None: return "Web search unavailable because beautifulsoup4 is not installed." @@ -70,6 +72,7 @@ def _call_openai( retries: int, previous_response_id: Optional[str] = None ) -> Tuple[dict, Optional[str]]: + _logging.debug(": Calling OpenAI API") payload_copy = payload.copy() if "input" not in payload_copy: payload_copy["input"] = str(input_data) @@ -207,6 +210,7 @@ def ai( description: Request summary text to be merged into the output. """ + _logging.info(f": Generating data using AI :: model :: {model}, thread_count :: {threads}, record_count :: {1 if not isinstance(input, list) else len(input)}") input_was_scalar = not isinstance(input, list) input_list = [input] if input_was_scalar else input diff --git a/wrangles/lookup.py b/wrangles/lookup.py index f2882cab5..49a6eb58d 100644 --- a/wrangles/lookup.py +++ b/wrangles/lookup.py @@ -1,4 +1,5 @@ from typing import Union as _Union +import logging as _logging from . import config as _config from . import data as _data from . import batching as _batching @@ -8,6 +9,7 @@ def lookup( input: _Union[str, list], model_id: str, columns: _Union[str, list] = None, + n: int = None, **kwargs ) -> _Union[str, list]: """ @@ -16,7 +18,9 @@ def lookup( :param input: A value or list of values to be looked up. :param model_id: The model to be used. :param columns: (Optional) The columns to be returned. If not provided, all columns will be returned as a dict. - """ + :param n: (Optional) Number of matches to return per input. When > 1, returns a list of n + dicts per input - each match is always a dict, even if a single column is requested. + """ # Check if user has entered a single input or multiple inputs single_input = False if not isinstance(input, list): @@ -56,6 +60,11 @@ def lookup( f'Using {purpose} model_id {model_id} in a lookup wrangle.' ) + _logging.info(f": Looking up {len(input)} values :: model_id :: {model_id}") + + if n: + kwargs['n'] = n + results = _batching.batch_api_calls( f'{_config.api_host}/wrangles/lookup', { @@ -67,20 +76,27 @@ def lookup( batch_size ) - if columns is None: - # If no columns specified, return as [{"col1": "val1", ...}, ...] - results = [ - {col: val for col, val in zip(results["columns"], row)} - for row in results["data"] - ] - elif single_columns: - # If single column specified, return as 1D array [val1, ...] - results = [r[0] for r in results["data"]] + if n and n > 1: + # API returns 1 row per input; row[0] is a list of n match dicts. + # When n > 1, every match is always returned as a dict, even if a + # single column was requested - naming an output column after a + # lookup column does not collapse the dict to that column's value. + results = [row[0] for row in results["data"]] else: - # If multiple columns specified, return as 2D array [[val1, ...], ...] - results = results["data"] + if columns is None: + # If no columns specified, return as [{"col1": "val1", ...}, ...] + results = [ + {col: val for col, val in zip(results["columns"], row)} + for row in results["data"] + ] + elif single_columns: + # If single column specified, return as 1D array [val1, ...] + results = [r[0] for r in results["data"]] + else: + # If multiple columns specified, return as 2D array [[val1, ...], ...] + results = results["data"] - # If input was a single value, return a single value + # If input was a single value, return a single value (or list of n matches) if single_input: results = results[0] diff --git a/wrangles/openai.py b/wrangles/openai.py index 32b303e25..898091e36 100644 --- a/wrangles/openai.py +++ b/wrangles/openai.py @@ -4,6 +4,7 @@ import copy as _copy import concurrent.futures as _futures from itertools import chain as _chain +import logging as _logging import requests as _requests import numpy as _np import time as _time @@ -52,9 +53,12 @@ def chatGPT( if not isinstance(retries, int) or retries < 0: raise ValueError("Retries must be a positive integer") - + + _logging.debug(f": Calling OpenAI ChatGPT :: timeout :: {timeout}, retries :: {retries}") + response = None backoff_time = 1 + retry_count = 0 while (retries + 1): try: response = _requests.post( @@ -100,9 +104,12 @@ def chatGPT( if "Incorrect API key" in error_message: raise ValueError("API Key provided is missing or invalid.") - retries -=1 - _time.sleep(backoff_time) - backoff_time *= 2 + retries -= 1 + retry_count += 1 + if retries >= 0: + _logging.warning(f": Retrying OpenAI request :: attempt :: {retry_count}") + _time.sleep(backoff_time) + backoff_time *= 2 if response and response.ok: try: @@ -117,7 +124,9 @@ def chatGPT( error_message = response.json()['error']['message'] except: error_message = "Failed" - + + _logging.error(f": OpenAI API error :: {error_message}") + # Return error for each requested column return { param: error_message @@ -155,6 +164,9 @@ def _embedding_thread( """ if request_params is None: request_params = {} + + _logging.debug(f": Computing embeddings :: model :: {model}, record_count :: {len(input_list)}") + response = None backoff_time = 1 while (retries + 1): diff --git a/wrangles/recipe.py b/wrangles/recipe.py index ce19c4a4e..09149c4b9 100644 --- a/wrangles/recipe.py +++ b/wrangles/recipe.py @@ -13,6 +13,7 @@ import re as _re import warnings as _warnings import concurrent.futures as _futures +import time as _time import pandas as _pandas import requests as _requests from . import recipe_wrangles as _recipe_wrangles @@ -63,6 +64,11 @@ def _load_recipe( """ if variables is None: variables = {} + + # Accept path-like objects (e.g. pathlib.Path) by converting to str + if isinstance(recipe, _os.PathLike): + recipe = str(recipe) + if isinstance(recipe, str) and "\n" not in recipe: _logging.info(f": Reading Recipe :: {recipe}") @@ -443,6 +449,27 @@ def _execute_wrangles( # Used to store parameters common to all wrangles - e.g where common_params = {} + # If the action is conditional, check if it should be run + # before column validation or wildcard expansion so that + # wrangles can be skipped when their input columns don't exist + if ( + "if" in params and + not _evaluate_conditional( + params["if"], + { + **variables, + **{ + "row_count": len(df), + "column_count": len(df.columns), + "columns": df.columns.tolist(), + "df": df + } + } + ) + ): + _logging.info(f": Wrangling :: {wrangle} skipped due to not passing the if statement.") + continue + # Blacklist of Wrangles not to allow wildcards for original_input = params.get('input') # Save for later reference if ( @@ -479,30 +506,50 @@ def _execute_wrangles( preserve_index=True ) - # If the action is conditional, check if it should be run - if ( - "if" in params and - not _evaluate_conditional( - params["if"], - { - **variables, - **{ - "row_count": len(df), - "column_count": len(df.columns), - "columns": df.columns.tolist(), - "df": df - } - } - ) - ): - _logging.info(f": Wrangling :: {wrangle} skipped due to not passing the if statement.") - continue + # If where filters out all rows, skip actually executing the + # wrangle (some wrangles error when given no rows), but if it + # explicitly declares output columns, still add them (empty) + # so the dataframe structure stays consistent with runs where + # at least one row matches - otherwise batching with where can + # produce inconsistent columns between batches + if len(df) == 0: + if 'output' in params: + if isinstance(params['output'], list): + output_columns = [ + list(col.values()) if isinstance(col, dict) else [col] + for col in params['output'] + ] + output_columns = [ + item + for sublist in output_columns + for item in sublist + ] + elif isinstance(params['output'], dict): + output_columns = list(params['output'].keys()) + else: + output_columns = [params['output']] + + for col in output_columns: + # Wildcard outputs (e.g. 'Col*') are expanded into + # concrete column names based on the actual data, + # which can't be determined with no rows to work + # with - skip adding those, only add named columns + if '*' in str(col): + continue + if col not in df_original.columns: + df_original[col] = '' + + df = df_original + continue # Add to common_params dict and remove from params for key in ['where', 'where_params', 'if']: if key in params.keys(): common_params[key] = params.pop(key) + _logging.info(f": Wrangling :: {wrangle} :: Starting") + _wrangle_start_time = _time.perf_counter() + if wrangle.split('.')[0] == 'pandas': # Execute a pandas method # TODO: disallow any hidden methods @@ -785,8 +832,9 @@ def _execute_wrangles( if isinstance(input_display, list): input_display = ', '.join(str(x) for x in input_display) - output_display = ', '.join(str(col) for col in output_columns) - _logging.info(f": Wrangling :: {wrangle} :: {input_display} >> {output_display}") + output_display = ', '.join(str(col) for col in output_columns) + _wrangle_elapsed = _time.perf_counter() - _wrangle_start_time + _logging.info(f": Wrangling :: {wrangle} :: Completed :: {input_display} >> {output_display} :: {_wrangle_elapsed:.3f}s") except Exception as e: # Append name of wrangle to message and pass through exception diff --git a/wrangles/recipe_wrangles/compare.py b/wrangles/recipe_wrangles/compare.py index c7b5e6940..879d09563 100644 --- a/wrangles/recipe_wrangles/compare.py +++ b/wrangles/recipe_wrangles/compare.py @@ -2,6 +2,7 @@ Functions to compare data from within columns """ +import logging as _logging import pandas as _pd from .. import compare as _compare @@ -43,6 +44,7 @@ def lists( description: Ignore case when comparing string items """ + _logging.debug(f": Comparing lists :: method :: {method}") if method not in ["intersection", "difference", "union"]: raise ValueError( "Method must be one of 'intersection', 'difference', 'union'" @@ -198,6 +200,7 @@ def text( description: "(Optional) Whether the comparison is case sensitive. Default is True" """ + _logging.debug(f": Comparing text strings :: input :: {input}") if method not in ["difference", "intersection", "overlap"]: raise ValueError( "Method must be one of 'overlap', 'difference' or 'intersection'" diff --git a/wrangles/recipe_wrangles/compute.py b/wrangles/recipe_wrangles/compute.py index f50400f9c..f06f53ba8 100644 --- a/wrangles/recipe_wrangles/compute.py +++ b/wrangles/recipe_wrangles/compute.py @@ -1,11 +1,15 @@ import pandas as _pd import numpy as _np +import re as _re +import logging as _logging + # Import our core compute and format functions from .. import compute as _compute from .. import format as _format + def case_when( df: _pd.DataFrame, output: str, @@ -44,6 +48,7 @@ def case_when( description: Value to assign if no conditions are met. Default None. """ + _logging.debug(f": Evaluating case_when :: condition_count :: {len(cases)}, output :: {output}") df_temp = df.copy() df_temp.columns = df_temp.columns.str.replace( r'[^a-zA-Z0-9_]', '_', regex=True) diff --git a/wrangles/recipe_wrangles/convert.py b/wrangles/recipe_wrangles/convert.py index 07fc12741..1c5cb9759 100644 --- a/wrangles/recipe_wrangles/convert.py +++ b/wrangles/recipe_wrangles/convert.py @@ -16,6 +16,14 @@ except ImportError: from yaml import SafeLoader as _YAMLLoader, SafeDumper as _YAMLDumper +# Pre-compiled regex for sentence case: matches the first non-whitespace character +# at the start of the string or immediately after punctuation (. ! ?) + optional +# whitespace. Using (\S) rather than ([a-zA-Z]) preserves the original char-by-char +# behaviour where a digit after punctuation "consumes" the capitalize flag +# (e.g. "13.5mm" stays lowercase because the `5` consumes the flag before `m` is +# reached), while \s* preserves handling of newlines and other Unicode whitespace. +_SENTENCE_CASE_RE = _re.compile(r'(^\s*|[.!?]\s*)(\S)') + def case(df: _pd.DataFrame, input: _Union[str, int, list], output: _Union[str, list] = None, case: str = 'lower') -> _pd.DataFrame: """ @@ -75,33 +83,44 @@ def case(df: _pd.DataFrame, input: _Union[str, int, list], output: _Union[str, l # Loop through and apply for all columns for input_column, output_column in zip(input, output): if desired_case != 'sentence': - df[output_column] = df[input_column].apply(lambda x: _safe_str_transform(x, desired_case, warnings)) + source = df[input_column] + # .str accessor raises AttributeError on non-object/non-string dtypes (e.g. int64). + # For those columns, preserve originals and warn once — matching prior behaviour. + if _pd.api.types.is_string_dtype(source.dtype) or _pd.api.types.is_object_dtype(source.dtype): + # Use vectorized pandas str methods (C-level, much faster than row-wise apply). + # Non-string values (lists, ints, etc.) become NaN after str operations; + # detect those and restore the originals so behaviour is unchanged. + transformed = getattr(source.str, desired_case)() + non_string_mask = source.notna() & transformed.isna() + if non_string_mask.any(): + if not warnings["invalid_data"]["logged"]: + _logging.warning(warnings['invalid_data']['message']) + warnings["invalid_data"]['logged'] = True + transformed = transformed.where(~non_string_mask, source) + df[output_column] = transformed + else: + if not warnings["invalid_data"]["logged"]: + _logging.warning(warnings['invalid_data']['message']) + warnings["invalid_data"]['logged'] = True + df[output_column] = source - elif desired_case == 'sentence': - def _getSentenceCase(source: str, warnings={}): + else: + # Sentence case: lowercase everything with the vectorized str method, then + # use a pre-compiled regex to re-capitalise sentence starts. This avoids + # the slow character-by-character Python loop of the previous implementation. + def _getSentenceCase(source, warnings=warnings): if isinstance(source, str): - output = [] - isFirstWord = True - - for character in source: - if isFirstWord and not character.isspace(): - output.append(character.upper()) - isFirstWord = False - elif not isFirstWord and character in ".!?": - isFirstWord = True - output.append(character.upper()) - else: - output.append(character.lower()) - - return ''.join(output) + return _SENTENCE_CASE_RE.sub( + lambda m: m.group(1) + m.group(2).upper(), + source.lower() + ) else: - # Only show this once to not spam the logs - if not warnings.get("invalid_data", {}).get('logged', False): + if not warnings["invalid_data"]["logged"]: _logging.warning(warnings['invalid_data']['message']) warnings["invalid_data"]['logged'] = True return source - df[output_column] = df[input_column].apply(lambda x: _getSentenceCase(x, warnings)) + df[output_column] = df[input_column].apply(_getSentenceCase) return df @@ -325,9 +344,30 @@ def fraction_to_decimal( return df +def _normalize_default_list(default, input: list, func_name: str) -> list: + """ + Normalize a default value/list into a per-column list of defaults. + + A scalar (or a single-element list) is broadcast to all columns. + A list with the same length as input is used as-is, one default per column. + An empty list (`[]`) is treated as the default value itself, not a per-column list. + """ + if isinstance(default, list) and len(default) > 0: + if len(default) == 1: + return [default[0]] * len(input) + elif len(default) != len(input): + raise ValueError( + f'The list of default values must be a single value or the same length as input/output for {func_name}' + ) + else: + return default + else: + return [default] * len(input) + + def from_json( - df: _pd.DataFrame, - input: _Union[str, int, list], + df: _pd.DataFrame, + input: _Union[str, int, list], output: _Union[str, list] = None, default = None, **kwargs @@ -351,43 +391,43 @@ def from_json( description: Name of the output column. If omitted, the input column will be overwritten default: type: ["string","array","object","number","boolean","null"] - description: Value to return if the row is empty or fails to be parsed as JSON + description: >- + Value to return if the row is empty or fails to be parsed as JSON. + If input is a list, default may also be a list - either a single + value to apply to all columns, or one value per input column. """ - def _load_with_fallback(value): - """ - Attempt to load JSON. - If fails and user has provided a default, return that. - If no default, raise an error. - """ - try: - return _json.loads(value, **kwargs) - except: - if default != None: - return default - else: - raise ValueError( - "Unable to load all rows as JSON. " + - "Set a default to set a value if the row is empty or fails to parse." - ) - # Set output column as input if not provided if output is None: output = input - + # Ensure input and outputs are lists if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] - + # Ensure input and output are equal lengths if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') - + + defaults = _normalize_default_list(default, input, 'convert.from_json') + # Loop through and apply for all columns - for input_column, output_column in zip(input, output): + for input_column, output_column, col_default in zip(input, output, defaults): + def _load_with_fallback(value, _col_default=col_default): + try: + return _json.loads(value, **kwargs) + except: + if _col_default is not None: + return _col_default + else: + raise ValueError( + "Unable to load all rows as JSON. " + + "Set a default to set a value if the row is empty or fails to parse." + ) + df[output_column] = [ _load_with_fallback(x) for x in df[input_column] ] - + return df @@ -507,43 +547,47 @@ def from_yaml( If omitted, the input column will be overwritten default: type: ["string","array","object","number","boolean","null"] - description: Value to return if the row is empty or fails to be parsed as JSON + description: >- + Value to return if the row is empty or fails to be parsed as YAML. + If input is a list, default may also be a list - either a single + value to apply to all columns, or one value per input column. """ - def _load_with_fallback(value): - """ - Attempt to load JSON. - If fails and user has provided a default, return that. - If no default, raise an error. - """ - try: - return _yaml.load(value, Loader=_YAMLLoader, **kwargs) or default - except: - if default != None: - return default - else: - raise ValueError( - "Unable to load all rows as YAML. " + - "Set a default to set a value if the row is empty or fails to parse." - ) - # Set output column as input if not provided if output is None: output = input - + # Ensure input and outputs are lists if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] - + # Ensure input and output are equal lengths if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') - + + defaults = _normalize_default_list(default, input, 'convert.from_yaml') + # Loop through and apply for all columns - for input_column, output_column in zip(input, output): + for input_column, output_column, col_default in zip(input, output, defaults): + def _load_with_fallback(value, _col_default=col_default): + try: + result = _yaml.load(value, Loader=_YAMLLoader, **kwargs) + # Only fall back to the default if the row was empty (parses to None). + # A falsy-but-valid result (e.g. [], {}, '', False, 0) must be preserved + # as-is rather than being overwritten by the default. + return _col_default if result is None else result + except: + if _col_default is not None: + return _col_default + else: + raise ValueError( + "Unable to load all rows as YAML. " + + "Set a default to set a value if the row is empty or fails to parse." + ) + df[output_column] = [ _load_with_fallback(x) for x in df[input_column] ] - + return df diff --git a/wrangles/recipe_wrangles/create.py b/wrangles/recipe_wrangles/create.py index ddda107dd..fdc755f3a 100644 --- a/wrangles/recipe_wrangles/create.py +++ b/wrangles/recipe_wrangles/create.py @@ -4,6 +4,7 @@ import uuid as _uuid from typing import Union as _Union import math as _math +import logging as _logging import pandas as _pd import numpy as _np import re as _re @@ -62,6 +63,7 @@ def bins( if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') + _logging.debug(f": Creating bins :: output :: {output}") for in_col, out_col in zip(input, output): # Dealing with positive infinity. At end of bins list if isinstance(bins, list): @@ -83,7 +85,13 @@ def bins( return df -def column(df: _pd.DataFrame, output: _Union[str, list], value = None) -> _pd.DataFrame: +def column( + df: _pd.DataFrame, + output: _Union[str, list], + value = None, + value_if_exists: str = 'existing', + coalesce_value: str = 'existing' +) -> _pd.DataFrame: """ type: object description: Create column(s) with a user defined value. Defaults to None (empty). @@ -104,12 +112,45 @@ def column(df: _pd.DataFrame, output: _Union[str, list], value = None) -> _pd.Da - array - boolean description: (Optional) Value(s) to add in the new column(s). If using a dictionary in output, value can only be a string. + value_if_exists: + type: string + description: >- + Determines behaviour when the output column already exists. + existing (default): leave the column unchanged. + coalesce: fill empty/null cells with the new value, keeping non-null cells. + new: overwrite the entire column with the new value. + enum: + - existing + - coalesce + - new + coalesce_value: + type: string + description: >- + Only used when value_if_exists is coalesce. Determines which side + is preferred when both the existing and new values are non-empty. + existing (default): keep the existing value, only fill empty/null cells with the new value. + new: keep the new value, only fall back to the existing value where the new value is empty/null. + enum: + - existing + - new """ + _logging.debug(f": Creating column(s) :: output :: {output}") + + valid_value_if_exists = ('existing', 'coalesce', 'new') + if value_if_exists not in valid_value_if_exists: + raise ValueError( + f'value_if_exists must be one of {valid_value_if_exists}, got "{value_if_exists}"' + ) + + valid_coalesce_value = ('existing', 'new') + if coalesce_value not in valid_coalesce_value: + raise ValueError( + f'coalesce_value must be one of {valid_coalesce_value}, got "{coalesce_value}"' + ) + # If a string provided, convert to list if isinstance(output, str): - if output in df.columns: - raise ValueError(f'"{output}" column already exists in dataFrame.') - output = [output] + output = [output] # gather the columns and values in a dictionary, if not a dict then use value as the value of dictionary output_dict = {} @@ -121,18 +162,36 @@ def column(df: _pd.DataFrame, output: _Union[str, list], value = None) -> _pd.Da else: output_dict.update({out: value}) - # Check if the list of outputs exist in dataFrame - check_list = [x for x in (output_dict.keys()) if x in df.columns] - if len(check_list) > 0: - raise ValueError(f'{check_list} column(s) already exists in the dataFrame') - for output_column, values_list in zip(output_dict.keys(), output_dict.values()): + column_exists = output_column in df.columns + + if column_exists and value_if_exists == 'existing': + continue + # Data to generate - data = _pd.DataFrame({ - output_column: _generate_cell_values(values_list, len(df)) - }).set_index(df.index) # use the same index as original to match rows - # Merging existing dataframe with values created - df = _pd.concat([df, data], axis=1) + new_values = _pd.Series( + _generate_cell_values(values_list, len(df)), + index=df.index + ) + + if column_exists and value_if_exists == 'coalesce': + if coalesce_value == 'existing': + primary, fallback = df[output_column], new_values + else: + primary, fallback = new_values, df[output_column] + + is_empty = primary.isna() + if primary.dtype == object: + is_empty = is_empty | (primary == '') + + df[output_column] = primary.where(~is_empty, fallback) + else: + # new or column doesn't exist — write/overwrite directly + if not column_exists: + data = _pd.DataFrame({output_column: new_values}) + df = _pd.concat([df, data], axis=1) + else: + df[output_column] = new_values return df @@ -219,6 +278,7 @@ def embeddings( if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') + _logging.info(f": Generating embeddings :: input_count :: {len(df)}, model :: {model}") if output_type not in ["python list", "numpy array"]: raise ValueError('Output_type must be of value "numpy array" or "python list"') @@ -258,6 +318,7 @@ def guid(df: _pd.DataFrame, output: _Union[str, list]) -> _pd.DataFrame: - array description: Name or list of names of new columns """ + _logging.debug(f": Creating GUIDs :: output :: {output}") return uuid(df, output) @@ -296,6 +357,7 @@ def index( if by != None and not isinstance(by, list): by = [by] + _logging.debug(f": Creating index column :: output :: {output}, start :: {start}") # If a string provided, convert to list if isinstance(output, str): output = [output] @@ -358,6 +420,7 @@ def jinja(df: _pd.DataFrame, template: dict, output: list, input: str = None) -> type: string description: A string which is used as the jinja template """ + _logging.debug(": Rendering Jinja template") if isinstance(output, list): output = output[0] @@ -417,6 +480,7 @@ def uuid(df: _pd.DataFrame, output: _Union[str, list]) -> _pd.DataFrame: - array description: Name or list of names of new columns """ + _logging.debug(f": Generating UUIDs :: output :: {output}") # If a string provided, convert to list if isinstance(output, str): output = [output] @@ -461,7 +525,8 @@ def hash(df: _pd.DataFrame, input: _Union[str, int, list], output: _Union[str, l if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') - + + _logging.debug(f": Hashing values :: method :: {method}") if method not in ['md5', 'sha1', 'sha256', 'sha512']: raise ValueError('Method must be one of: md5, sha1, sha256, sha512') diff --git a/wrangles/recipe_wrangles/extract.py b/wrangles/recipe_wrangles/extract.py index 6bb76bbd3..be825ee6e 100644 --- a/wrangles/recipe_wrangles/extract.py +++ b/wrangles/recipe_wrangles/extract.py @@ -3,11 +3,182 @@ """ from typing import Union as _Union import re as _re +import logging as _logging import pandas as _pd from .. import extract as _extract from .. import data as _data +_OUTPUT_FORMAT_ALIASES = { + "json": "json", + "json list": "json_list", + "json_list": "json_list", + "list": "json_list", + "array": "json_list", + "json dictionary": "json_dictionary", + "json dict": "json_dictionary", + "json_dictionary": "json_dictionary", + "json_dict": "json_dictionary", + "dict": "json_dictionary", + "dictionary": "json_dictionary", + "columns": "columns", + "column": "columns", + "concatenate": "concatenate", + "concat": "concatenate", +} + + +def _normalize_output_format(output_format, default): + if output_format is None: + return default + + output_format = _OUTPUT_FORMAT_ALIASES.get( + str(output_format).strip().lower().replace("-", "_"), + output_format + ) + + if output_format == "json": + return default + + if output_format not in ("json_list", "json_dictionary", "columns", "concatenate"): + raise ValueError( + "output_format must be one of list, dictionary, columns, or concatenate" + ) + + return output_format + + +def _ensure_list(value): + return value if isinstance(value, list) else [value] + + +def _is_columns_target(output, output_format, output_is_list=False): + """ + Whether results should be split across explicit output columns. + + An explicit output_format always wins. Otherwise, providing output + as a list of one or more column names implies Columns format. + """ + if output_format is not None: + return _normalize_output_format(output_format, "json_list") == "columns" + return output_is_list or (isinstance(output, list) and len(output) > 1) + + +def _resolve_output_format(output, output_format, default_format, output_is_list=False): + if output_format is None and (output_is_list or (isinstance(output, list) and len(output) > 1)): + return "columns" + return _normalize_output_format(output_format, default_format) + + +def _stringify_list(value, char): + if value in (None, ""): + return "" + if isinstance(value, list): + return char.join([str(item) for item in value]) + return str(value) + + +def _write_list_output( + df, + output, + results, + output_format, + char=", ", + default_format="json_list", + output_is_list=False +): + output_format = _resolve_output_format(output, output_format, default_format, output_is_list) + + if output_format == "json_dictionary": + raise ValueError("output_format dictionary is only valid for dictionary-producing extracts") + + if output_format == "concatenate": + df[output[0]] = [_stringify_list(row, char) for row in results] + return + + if output_format == "columns": + # A scalar (non-list) result counts as a single match. + # Always create at least one column, even if no rows + # produced any results + max_result_len = max( + [len(row) if isinstance(row, list) else 1 for row in results] + [1] + ) + if len(output) > 1 or output_is_list: + # Cap the number of columns created to however many + # output column names were explicitly provided, dropping + # any results beyond that even if more were found + output_columns = output[:min(len(output), max_result_len)] + else: + # No explicit column names given (a bare string output) + # with Columns format requested - auto-number columns + # for however many results were found + output_columns = [f"{output[0]} {i + 1}" for i in range(max_result_len)] + for i, output_column in enumerate(output_columns): + df[output_column] = [ + (row[i] if len(row) > i else "") if isinstance(row, list) + else (row if i == 0 else "") + for row in results + ] + return + + df[output[0]] = results + + +def _dict_keys(results): + keys = [] + for row in results: + if isinstance(row, dict): + for key in row: + if key not in keys: + keys.append(key) + return keys + + +def _write_dict_output(df, output, results, output_format, default_format="json_dictionary", output_is_list=False): + output_format = _resolve_output_format(output, output_format, default_format, output_is_list) + + if output_format in ("json_list", "concatenate"): + raise ValueError("output_format list or concatenate is only valid for list-producing extracts") + + if output_format == "columns": + output_columns = ( + output + if len(output) > 1 or output_is_list + else (_dict_keys(results) or output) + ) + for output_column in output_columns: + df[output_column] = [ + row.get(output_column, "") if isinstance(row, dict) else "" + for row in results + ] + return + + df[output[0]] = results + + +def _write_results( + df, + output, + results, + output_format, + char=", ", + default_format="json_list", + output_is_list=False +): + if default_format == "json_dictionary": + _write_dict_output(df, output, results, output_format, default_format, output_is_list) + else: + _write_list_output( + df, + output, + results, + output_format, + char, + default_format, + output_is_list + ) + + def _combine_list_rows(rows): combined = [] for row in rows: @@ -18,11 +189,24 @@ def _combine_list_rows(rows): return list(dict.fromkeys(combined)) +def _combine_dict_rows(rows): + combined = {} + for row in rows: + if not isinstance(row, dict): + continue + for key, value in row.items(): + values = value if isinstance(value, list) else [value] + combined[key] = _combine_list_rows([combined.get(key, []), values]) + return combined + + def address( df: _pd.DataFrame, input: _Union[str, int, list], output: _Union[str, list], dataType: str, + output_format: str = None, + char: str = ", ", **kwargs ) -> _pd.DataFrame: """ @@ -51,33 +235,55 @@ def address( - cities - regions - countries + output_format: + type: string + description: Format of the extract output + enum: + - list + - columns + - concatenate + char: + type: string + description: Character to use when output_format is concatenate """ # If output is not specified, overwrite input columns in place if output is None: output = input + # Whether output was explicitly given as a list of column names + output_is_list = isinstance(output, list) and len(output) > 1 + # If a string provided, convert to list if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] # Ensure input and output lengths are compatible - if len(input) != len(output) and len(output) > 1: + if len(input) != len(output) and len(output) > 1 and not (len(input) == 1 and _is_columns_target(output, output_format, output_is_list)): raise ValueError('Extract must output to a single column or equal amount of columns as input.') - if len(output) == 1 and len(input) > 1: - df[output[0]] = _extract.address( + if len(input) == 1 and _is_columns_target(output, output_format, output_is_list): + results = _extract.address( + df[input[0]].astype(str).tolist(), + dataType, + **kwargs + ) + _write_list_output(df, output, results, output_format, char, output_is_list=output_is_list) + elif len(output) == 1 and len(input) > 1: + results = _extract.address( df[input].astype(str).aggregate(' '.join, axis=1).tolist(), dataType, **kwargs ) + _write_list_output(df, output, results, output_format, char, output_is_list=output_is_list) else: # Loop through and apply for all columns for input_column, output_column in zip(input, output): - df[output_column] = _extract.address( + results = _extract.address( df[input_column].astype(str).tolist(), dataType, **kwargs ) - + _write_list_output(df, [output_column], results, output_format, char) + return df @@ -87,6 +293,8 @@ def ai( input: list = None, output: _Union[dict, str, list] = None, model_id: str = None, + output_format: str = None, + char: str = ", ", **kwargs ): """ @@ -179,7 +387,19 @@ def ai( Enable strict mode. Default False. If True, the function will be required to match the schema, but may be more limited in the schema it can return. + output_format: + type: string + description: Format of the extract output + enum: + - dictionary + - columns + - concatenate + char: + type: string + description: Character to use when output_format is concatenate """ + output_format_normalized = _normalize_output_format(output_format, "columns") + # If input is provided, extract only those columns # Otherwise, provide the whole dataframe if input is not None: @@ -251,7 +471,24 @@ def ai( try: exploded_df = _pd.json_normalize(results, max_level=0).fillna('').set_index(df.index) - if target_columns and len(target_columns) == 1: + if output_format_normalized == "json_dictionary": + output_column_name = target_columns[0] if target_columns and len(target_columns) == 1 else "output" + df[output_column_name] = results + elif output_format_normalized == "concatenate": + if target_columns and len(target_columns) != 1: + raise ValueError("output_format concatenate can only be used with a single output column") + output_column = target_columns[0] if target_columns else "output" + if len(exploded_df.columns) == 1: + df[output_column] = [ + _stringify_list(row, char) + for row in exploded_df[exploded_df.columns[0]].tolist() + ] + else: + df[output_column] = [ + char.join([_stringify_list(value, char) for value in row.values()]) + for row in results + ] + elif target_columns and len(target_columns) == 1: if len(exploded_df.columns) == 1: # If the AI model only returns a single column # then use the contents of that columns as the output @@ -286,6 +523,8 @@ def attributes( desired_unit: str = None, bound: str = 'mid', first_element: bool = False, + output_format: str = None, + char: str = ", ", **kwargs ) -> _pd.DataFrame: """ @@ -354,43 +593,93 @@ def attributes( first_element: type: boolean description: Get the first element from results + output_format: + type: string + description: Format of the extract output + enum: + - list + - dictionary + - columns + - concatenate + char: + type: string + description: Character to use when output_format is concatenate $ref: "#/$defs/misc/unit_entity_map" """ # If output is not specified, overwrite input columns in place if output is None: output = input + # Whether output was explicitly given as a list of column names + output_is_list = isinstance(output, list) and len(output) > 1 + # If a string provided, convert to list if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] # Ensure input and output lengths are compatible - if len(input) != len(output) and len(output) > 1: + if len(input) != len(output) and len(output) > 1 and not (len(input) == 1 and _is_columns_target(output, output_format, output_is_list)): raise ValueError('Extract must output to a single column or equal amount of columns as input.') - - if len(output) == 1 and len(input) > 1: + + if len(input) == 1 and _is_columns_target(output, output_format, output_is_list): + results = _extract.attributes( + df[input[0]].astype(str).tolist(), + responseContent, + attribute_type, + desired_unit, + bound, + False, + **kwargs + ) + _write_results( + df, + output, + results, + output_format, + char, + "json_list" if attribute_type else "json_dictionary", + output_is_list + ) + elif len(output) == 1 and len(input) > 1: # df[output[0]] = _extract.attributes(df[input].astype(str).aggregate(' AAA '.join, axis=1).tolist()) - df[output[0]] = _extract.attributes( + results = _extract.attributes( df[input].astype(str).aggregate(' AAA '.join, axis=1).tolist(), responseContent, attribute_type, desired_unit, bound, - first_element, + first_element if output_format is None else False, **kwargs ) + _write_results( + df, + output, + results, + output_format, + char, + "json_list" if attribute_type else "json_dictionary", + output_is_list + ) else: # Loop through and apply for all columns for input_column, output_column in zip(input, output): - df[output_column] = _extract.attributes( + results = _extract.attributes( df[input_column].astype(str).tolist(), responseContent, attribute_type, desired_unit, bound, - first_element, + first_element if output_format is None else False, **kwargs ) - + _write_results( + df, + [output_column], + results, + output_format, + char, + "json_list" if attribute_type else "json_dictionary" + ) + return df @@ -399,7 +688,9 @@ def brackets( input: _Union[str, int, list], output: _Union[str, list], find: _Union[str, list] = 'all', - include_brackets: bool = False + include_brackets: bool = False, + output_format: str = None, + char: str = ", " ) -> _pd.DataFrame: """ type: object @@ -428,18 +719,32 @@ def brackets( include_brackets: type: boolean description: (Optional) Include the brackets in the output + output_format: + type: string + description: Format of the extract output + enum: + - list + - columns + - concatenate + char: + type: string + description: Character to use when output_format is concatenate """ # If output is not specified, overwrite input columns in place if output is None: output = input + # Whether output was explicitly given as a list of column names + output_is_list = isinstance(output, list) and len(output) > 1 + # If a string provided, convert to list if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] # Ensure input and output lengths are compatible - if len(input) != len(output) and len(output) > 1: + if len(input) != len(output) and len(output) > 1 and not (len(input) == 1 and _is_columns_target(output, output_format, output_is_list)): raise ValueError('Extract must output to a single column or equal amount of columns as input.') + _logging.debug(f": Extracting from brackets :: input :: {input}") # Ensure find is a list if not isinstance(find, list): find = [find] @@ -449,20 +754,51 @@ def brackets( if not all(element in bracket_types for element in find): raise ValueError("find must only contain the elements: round, square, curly, angled") - # If only only one output and multiple inputs, concatenate the inputs - if len(output) == 1 and len(input) > 1: - df[output[0]] = _extract.brackets( + if len(input) == 1 and _is_columns_target(output, output_format, output_is_list): + results = _extract.brackets( + df[input[0]].astype(str).tolist(), + find, + include_brackets, + return_data_type="list" + ) + _write_list_output( + df, + output, + results, + output_format, + char, + output_is_list=output_is_list + ) + elif len(output) == 1 and len(input) > 1: + results = _extract.brackets( df[input].astype(str).aggregate(' '.join, axis=1).tolist(), find, - include_brackets + include_brackets, + return_data_type="list" + ) + _write_list_output( + df, + output, + results, + output_format, + char, + output_is_list=output_is_list ) else: # Loop through and apply for all columns for input_column, output_column in zip(input, output): - df[output_column] = _extract.brackets( + results = _extract.brackets( df[input_column].astype(str).tolist(), find, - include_brackets + include_brackets, + return_data_type="list" + ) + _write_list_output( + df, + [output_column], + results, + output_format, + char ) return df @@ -473,6 +809,8 @@ def codes( input: _Union[str, int, list], output: _Union[str, list], first_element: bool = False, + output_format: str = None, + char: str = ", ", **kwargs ) -> _pd.DataFrame: """ @@ -496,6 +834,16 @@ def codes( first_element: type: boolean description: Get the first element from results + output_format: + type: string + description: Format of the extract output + enum: + - list + - columns + - concatenate + char: + type: string + description: Character to use when output_format is concatenate min_length: type: - integer @@ -515,8 +863,9 @@ def codes( - strict sort_order: type: string - description: Default is as found in the input. Also allows longest or shortest. + description: Default is input order. Also allows longest or shortest. enum: + - input - longest - shortest disallowed_patterns: @@ -525,31 +874,47 @@ def codes( include_multi_part_tokens: type: boolean description: Whether to include multi-part tokens that have a space. Default True. + extract_raw: + type: boolean + description: Whether to return tokens with their adjacent non-whitespace characters. Default False. """ # If output is not specified, overwrite input columns in place if output is None: output = input + # Whether output was explicitly given as a list of column names + output_is_list = isinstance(output, list) and len(output) > 1 + # If a string provided, convert to list if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] # Ensure input and output lengths are compatible - if len(input) != len(output) and len(output) > 1: + if len(input) != len(output) and len(output) > 1 and not (len(input) == 1 and _is_columns_target(output, output_format, output_is_list)): raise ValueError('Extract must output to a single column or equal amount of columns as input.') - if len(output) == 1 and len(input) > 1: - df[output[0]] = _extract.codes( + if len(input) == 1 and _is_columns_target(output, output_format, output_is_list): + results = _extract.codes( + df[input[0]].astype(str).tolist(), + False, + **kwargs + ) + _write_list_output(df, output, results, output_format, char, output_is_list=output_is_list) + elif len(output) == 1 and len(input) > 1: + results = _extract.codes( df[input].astype(str).aggregate(' AAA '.join, axis=1).tolist(), + first_element if output_format is None else False, **kwargs ) + _write_list_output(df, output, results, output_format, char, output_is_list=output_is_list) else: # Loop through and apply for all columns for input_column, output_column in zip(input, output): - df[output_column] = _extract.codes( + results = _extract.codes( df[input_column].astype(str).tolist(), - first_element, + first_element if output_format is None else False, **kwargs ) + _write_list_output(df, [output_column], results, output_format, char) return df @@ -564,7 +929,10 @@ def custom( case_sensitive: bool = False, extract_raw: bool = False, use_spellcheck: bool = False, + include_empty_labels: bool = True, sort: str = 'training_order', + output_format: str = None, + char: str = ", ", **kwargs ) -> _pd.DataFrame: """ @@ -617,34 +985,69 @@ def custom( - reverse_alphabetical - ascending - descending + output_format: + type: string + description: Format of the extract output + enum: + - list + - dictionary + - columns + - concatenate + char: + type: string + description: Character to use when output_format is concatenate + include_empty_labels: + type: boolean + description: Include labels with no found values in the output when using use_labels=True """ if output is None: output = input - + + # Whether output was explicitly given as a list of column names + output_is_list = isinstance(output, list) and len(output) > 1 + # If a string provided, convert to list if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] if not isinstance(model_id, list): model_id = [model_id] - # Ensure input and output lengths are compatible - if len(input) != len(output) and len(output) > 1: + default_format = "json_dictionary" if use_labels else "json_list" + + if len(input) != len(output) and len(output) > 1 and not (len(input) == 1 and _is_columns_target(output, output_format, output_is_list)): raise ValueError('Extract must output to a single column or equal amount of columns as input.') - if len(input) == len(output) and len(model_id) == 1: + if len(input) == 1 and _is_columns_target(output, output_format, output_is_list) and len(model_id) == 1: + results = _extract.custom( + df[input[0]].astype(str).tolist(), + model_id=model_id[0], + first_element=False, + use_labels=use_labels, + case_sensitive=case_sensitive, + extract_raw=extract_raw, + use_spellcheck=use_spellcheck, + include_empty_labels=include_empty_labels, + sort=sort, + **kwargs + ) + _write_results(df, output, results, output_format, char, default_format, output_is_list) + + elif len(input) == len(output) and len(model_id) == 1: # if one model_id, then use that model for all columns inputs and outputs model_id = [model_id[0] for _ in range(len(input))] for in_col, out_col, model in zip(input, output, model_id): - df[out_col] = _extract.custom( + results = _extract.custom( df[in_col].astype(str).tolist(), model_id=model, - first_element=first_element, + first_element=first_element if output_format is None else False, use_labels=use_labels, case_sensitive=case_sensitive, extract_raw=extract_raw, use_spellcheck=use_spellcheck, + include_empty_labels=include_empty_labels, sort=sort, **kwargs ) - + _write_results(df, [out_col], results, output_format, char, default_format) + elif len(input) > 1 and len(output) == 1 and len(model_id) == 1: model_id = [model_id[0] for _ in range(len(input))] output = output[0] @@ -654,32 +1057,38 @@ def custom( df_temp[output + str(i)] = _extract.custom( df[in_col].astype(str).tolist(), model_id=single_model_id, - first_element=first_element, + first_element=first_element if output_format is None else False, use_labels=use_labels, case_sensitive=case_sensitive, extract_raw=extract_raw, use_spellcheck=use_spellcheck, + include_empty_labels=include_empty_labels, sort=sort, **kwargs ) - # Combine the per-column match lists into a single output list. - df[output] = [_combine_list_rows(row) for row in df_temp.values.tolist()] + if use_labels: + results = [_combine_dict_rows(row) for row in df_temp.values.tolist()] + else: + results = [_combine_list_rows(row) for row in df_temp.values.tolist()] + _write_results(df, [output], results, output_format, char, default_format, output_is_list) else: # Iterate through the inputs, outputs and model_ids for in_col, out_col, model in zip(input, output, model_id): - df[out_col] = _extract.custom( + results = _extract.custom( df[in_col].astype(str).tolist(), model_id=model, - first_element=first_element, + first_element=first_element if output_format is None else False, use_labels=use_labels, case_sensitive=case_sensitive, extract_raw=extract_raw, use_spellcheck=use_spellcheck, + include_empty_labels=include_empty_labels, sort=sort, **kwargs ) + _write_results(df, [out_col], results, output_format, char, default_format) return df @@ -727,7 +1136,8 @@ def date_properties(df: _pd.DataFrame, input: _pd.Timestamp, property: str, outp # Ensure input and output lengths are compatible if len(input) != len(output) and len(output) > 1: raise ValueError('Extract must output to a single column or equal amount of columns as input.') - + + _logging.debug(f": Extracting date property :: {property} from {input}") if len(output) == 1 and len(input) > 1: output = [output[0] for i in range(len(input))] # df_temp = df[input].apply(_pd.to_datetime) @@ -825,6 +1235,7 @@ def date_range(df: _pd.DataFrame, start_time: _pd.Timestamp, end_time: _pd.Times - seconds - milliseconds """ + _logging.debug(f": Generating date range :: output :: {output}") range_object = { 'business days': 'B', 'days': 'D', @@ -871,6 +1282,8 @@ def html( input: _Union[str, int, list], data_type: str, output: _Union[str, list] = None, + output_format: str = None, + char: str = ", ", **kwargs ) -> _pd.DataFrame: """ @@ -898,29 +1311,50 @@ def html( enum: - text - links - first_element: - type: boolean - description: Get the first element from results + output_format: + type: string + description: Format of the extract output + enum: + - list + - columns + - concatenate + char: + type: string + description: Character to use when output_format is concatenate """ # If output is not specified, overwrite input columns in place if output is None: output = input + # Whether output was explicitly given as a list of column names + output_is_list = isinstance(output, list) and len(output) > 1 + # If a string provided, convert to list if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] # Ensure input and output lengths are compatible - if len(input) != len(output) and len(output) > 1: + if len(input) != len(output) and len(output) > 1 and not (len(input) == 1 and _is_columns_target(output, output_format, output_is_list)): raise ValueError('Extract must output to a single column or equal amount of columns as input.') - - # Loop through and apply for all columns - for input_column, output_column in zip(input, output): - df[output_column] = _extract.html( - df[input_column].astype(str).tolist(), + + _logging.debug(f": Extracting from HTML :: input :: {input}") + + if len(input) == 1 and _is_columns_target(output, output_format, output_is_list): + results = _extract.html( + df[input[0]].astype(str).tolist(), dataType=data_type, **kwargs ) - + _write_list_output(df, output, results, output_format, char, output_is_list=output_is_list) + else: + # Loop through and apply for all columns + for input_column, output_column in zip(input, output): + results = _extract.html( + df[input_column].astype(str).tolist(), + dataType=data_type, + **kwargs + ) + _write_list_output(df, [output_column], results, output_format, char) + return df @@ -931,6 +1365,8 @@ def properties( property_type: str = None, return_data_type: str = 'list', first_element: bool = False, + output_format: str = None, + char: str = ", ", **kwargs ) -> _pd.DataFrame: """ @@ -961,44 +1397,95 @@ def properties( - Standards return_data_type: type: string - description: The format to return the data, as a list or as a string + description: Legacy format option. Prefer output_format. enum: - list - string first_element: type: boolean description: Get the first element from results + output_format: + type: string + description: Format of the extract output + enum: + - list + - dictionary + - columns + - concatenate + char: + type: string + description: Character to use when output_format is concatenate """ # If output is not specified, overwrite input columns in place if output is None: output = input + # Whether output was explicitly given as a list of column names + output_is_list = isinstance(output, list) and len(output) > 1 + # If a string provided, convert to list if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] # Ensure input and output lengths are compatible - if len(input) != len(output) and len(output) > 1: + if output_format is None and return_data_type == "string": + output_format = "concatenate" + + if len(input) != len(output) and len(output) > 1 and not (len(input) == 1 and _is_columns_target(output, output_format, output_is_list)): raise ValueError('Extract must output to a single column or equal amount of columns as input.') - if len(output) == 1 and len(input) > 1: - df[output[0]] = _extract.properties( + if len(input) == 1 and _is_columns_target(output, output_format, output_is_list): + results = _extract.properties( + df[input[0]].astype(str).tolist(), + type=property_type, + return_data_type='list', + first_element=False, + **kwargs + ) + _write_results( + df, + output, + results, + output_format, + char, + "json_list" if property_type else "json_dictionary", + output_is_list + ) + elif len(output) == 1 and len(input) > 1: + results = _extract.properties( df[input].astype(str).aggregate(' '.join, axis=1).tolist(), type=property_type, - return_data_type=return_data_type, - first_element=first_element, + return_data_type='list' if (output_format is not None or output_is_list) else return_data_type, + first_element=first_element if (output_format is None and not output_is_list) else False, **kwargs ) + _write_results( + df, + output, + results, + output_format, + char, + "json_list" if property_type else "json_dictionary", + output_is_list + ) else: # Loop through and apply for all columns for input_column, output_column in zip(input, output): - df[output_column] = _extract.properties( + results = _extract.properties( df[input_column].astype(str).tolist(), type=property_type, - return_data_type=return_data_type, - first_element=first_element, + return_data_type='list' if output_format is not None else return_data_type, + first_element=first_element if output_format is None else False, **kwargs ) - + _write_results( + df, + [output_column], + results, + output_format, + char, + "json_list" if property_type else "json_dictionary" + ) + return df def regex( @@ -1007,7 +1494,9 @@ def regex( find: str, output: _Union[str, list], output_pattern: str = None, - first_element: bool = False + first_element: bool = False, + output_format: str = None, + char: str = ", " ) -> _pd.DataFrame: r""" type: object @@ -1041,47 +1530,57 @@ def regex( first_element: type: boolean description: Get the first element from results + output_format: + type: string + description: Format of the extract output + enum: + - list + - columns + - concatenate + char: + type: string + description: Character to use when output_format is concatenate """ # If output is not specified, overwrite input columns in place - if output is None: + if output is None: output = input + # Whether output was explicitly given as a list of column names + output_is_list = isinstance(output, list) and len(output) > 1 + # If a string is provided, convert to list - if not isinstance(input, list): + if not isinstance(input, list): input = [input] - if not isinstance(output, list): + if not isinstance(output, list): output = [output] # Ensure input and output lengths are compatible - if len(input) != len(output) and len(output) > 1: + if len(input) != len(output) and len(output) > 1 and not (len(input) == 1 and _is_columns_target(output, output_format, output_is_list)): raise ValueError('Extract must output to a single column or equal amount of columns as input.') - + + _logging.debug(f": Extracting regex patterns :: input :: {input}") find_pattern = _re.compile(find) - - # Loop through and apply for all columns - for input_column, output_column in zip(input, output): - if output_pattern is None and first_element: - # Return entire matches - df[output_column] = df[input_column].apply( - lambda x: ([match.group(0) for match in _re.finditer(find_pattern, str(x) if x is not None else "")][0] - if len([match.group(0) for match in _re.finditer(find_pattern, str(x) if x is not None else "")]) >= 1 - else "") - ) - elif output_pattern is None and not first_element: - # Return entire matches - df[output_column] = df[input_column].apply(lambda x: [match.group(0) for match in _re.finditer(find_pattern, str(x) if x is not None else "")]) - elif output_pattern and first_element: - # Return specific capture groups in the pattern the were passed - df[output_column] = df[input_column].apply( - lambda x: ( - [find_pattern.sub(output_pattern, match.group(0)) for match in find_pattern.finditer(str(x) if x is not None else "")][0] - if len([find_pattern.sub(output_pattern, match.group(0)) for match in find_pattern.finditer(str(x) if x is not None else "")]) >= 1 - else "" - ) - ) + + def _matches(value): + value = str(value) if value is not None else "" + matches = [match.group(0) for match in _re.finditer(find_pattern, value)] + if output_pattern: + matches = [find_pattern.sub(output_pattern, match) for match in matches] + return matches + + def _write_regex(input_column, output_columns, columns_is_list=False): + results = df[input_column].apply(_matches).tolist() + if output_format is None and first_element and len(output_columns) == 1 and not columns_is_list: + df[output_columns[0]] = [row[0] if len(row) >= 1 else "" for row in results] else: - # Return specific capture groups in the pattern the were passed - df[output_column] = df[input_column].apply(lambda x: [find_pattern.sub(output_pattern, match.group(0)) for match in find_pattern.finditer(str(x) if x is not None else "")]) + _write_list_output(df, output_columns, results, output_format, char, output_is_list=columns_is_list) + + if len(input) == 1 and _is_columns_target(output, output_format, output_is_list): + _write_regex(input[0], output, output_is_list) + else: + # Loop through and apply for all columns + for input_column, output_column in zip(input, output): + _write_regex(input_column, [output_column]) return df diff --git a/wrangles/recipe_wrangles/format.py b/wrangles/recipe_wrangles/format.py index 2a8b64288..70c59c209 100644 --- a/wrangles/recipe_wrangles/format.py +++ b/wrangles/recipe_wrangles/format.py @@ -2,6 +2,7 @@ Functions to re-format data """ from typing import Union as _Union +import logging as _logging import pandas as _pd from .. import format as _format from ..utils import safe_str_transform as _safe_str_transform @@ -43,6 +44,7 @@ def dates(df: _pd.DataFrame, input: _Union[str, int, list], format: str, output: if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') + _logging.debug(f": Formatting dates :: format :: {format}, input :: {input}") # Loop through and apply for all columns for input_column, output_column in zip(input, output): # convert the column to timestamp type and format date @@ -98,6 +100,7 @@ def pad( description: If true, skip padding for empty or whitespace-only values default: false """ + _logging.debug(f": Padding strings :: pad_length :: {pad_length}, side :: {side}") char = str(char) # If the output is not specified, overwrite input columns in place if output is None: output = input @@ -171,7 +174,8 @@ def prefix( # If the input and output are not the same type if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') - + + _logging.debug(f": Adding prefix to {input}") # Loop through and apply for all columns for input_column, output_column in zip(input, output): if skip_empty: @@ -216,6 +220,7 @@ def remove_duplicates(df: _pd.DataFrame, input: _Union[str, int, list], output: if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') + _logging.debug(f": Removing duplicates :: input :: {input}, ignore_case :: {ignore_case}") # Loop through and apply for all columns for input_column, output_column in zip(input, output): df[output_column] = _format.remove_duplicates(df[input_column].values.tolist(), ignore_case) @@ -257,6 +262,7 @@ def significant_figures(df: _pd.DataFrame, input: _Union[str, int, list], signif if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') + _logging.debug(f": Rounding to {significant_figures} significant figures :: input :: {input}") # Loop through all requested columns and apply sig figs for input_column, output_column in zip(input, output): df[output_column] = _format.significant_figures(df[input_column].to_list(), significant_figures) @@ -311,6 +317,7 @@ def suffix( if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') + _logging.debug(f": Adding suffix to {input}") # Loop through and apply for all columns for input_column, output_column in zip(input, output): if skip_empty: @@ -351,6 +358,7 @@ def trim(df: _pd.DataFrame, input: _Union[str, int, list], output: _Union[str, l if len(input) != len(output): raise ValueError('The lists for input and output must be the same length.') + _logging.debug(f": Trimming whitespace :: input :: {input}") warnings = { "invalid_data": { "logged": False, @@ -370,5 +378,6 @@ def price_breaks(df: _pd.DataFrame, input: list, categoryLabel: str, valueLabel: """ Rearrange price breaks """ + _logging.info(f": Processing price breaks :: input :: {input}") df = _pd.concat([df, _format.price_breaks(df[input], categoryLabel, valueLabel)], axis=1) return df diff --git a/wrangles/recipe_wrangles/generate.py b/wrangles/recipe_wrangles/generate.py index 1c78af7a9..a4daaca04 100644 --- a/wrangles/recipe_wrangles/generate.py +++ b/wrangles/recipe_wrangles/generate.py @@ -1,5 +1,6 @@ from typing import Union as _Union, Dict as _Dict, List as _List, Optional as _Optional +import logging as _logging import pandas as _pd import wrangles.generate as _generate @@ -79,6 +80,7 @@ def ai( type: boolean description: Request summary text to be merged into the output. """ + _logging.info(f": Generating AI output :: model :: {model}, thread_count :: {threads}") if input is not None: if not isinstance(input, list): input = [input] diff --git a/wrangles/recipe_wrangles/main.py b/wrangles/recipe_wrangles/main.py index 2fa67d4b4..83426351b 100644 --- a/wrangles/recipe_wrangles/main.py +++ b/wrangles/recipe_wrangles/main.py @@ -903,7 +903,8 @@ def lookup( input: str, output: _Union[str, list] = None, model_id: str = None, - lookup_mode: str = 'by_row', + lookup_mode: str = 'by_row', + n: int = None, **kwargs ) -> _pd.DataFrame: """ @@ -925,7 +926,17 @@ def lookup( type: - string - array - description: Name of the output column(s) + description: >- + Name of the output column(s). When n is provided and the output list + length equals n, each output column receives the corresponding match. + A single output containing a wildcard (*) is expanded into n columns, + e.g. "Top *" with n: 3 becomes "Top 1", "Top 2", "Top 3". + n: + type: integer + description: >- + Number of matches to return per input value. When the output list + length equals n, each output column receives the corresponding match. + Otherwise all n matches are stored as a list in each output column. lookup_mode: type: string description: >- @@ -950,6 +961,16 @@ def lookup( # Ensure output is a list if not isinstance(output, list): output = [output] + # Expand a single wildcard output name into one column per match + # e.g. output: "Top *" with n=3 -> ["Top 1", "Top 2", "Top 3"] + if ( + n and n > 1 and + len(output) == 1 and + isinstance(output[0], str) and + '*' in output[0] + ): + output = [output[0].replace('*', str(i)) for i in range(1, n + 1)] + # Return early on empty df if df.empty: # Add empty output columns to maintain expected structure @@ -986,6 +1007,12 @@ def _clean_kwargs(kwargs): kwargs.pop('matrix_variables') return kwargs + # Distribute the n matches for each row across the output columns, + # ranked match i goes to output column i + def _distribute_n_matches(data): + for i, out in enumerate(output): + df[out] = [row[i] if isinstance(row, list) and i < len(row) else None for row in data] + # Perform lookup based on lookup_mode if lookup_mode == 'by_row': # Current behavior - process all rows @@ -995,20 +1022,40 @@ def _clean_kwargs(kwargs): df[input].values.tolist(), model_id, columns=wrangle_output, + n=n, **_clean_kwargs(kwargs) ) - df[output] = data + if n and n > 1 and len(output) == n: + # Distribute: each output column gets the nth match + _distribute_n_matches(data) + elif n and n > 1 and len(output) > 1: + raise ValueError( + f'When n > 1 and multiple output columns are provided, the number ' + f'of output columns ({len(output)}) must equal n ({n}).' + ) + else: + df[output] = data elif not any([col in metadata["settings"]["columns"] for col in wrangle_output]): # User specified no columns from the wrangle data = _lookup( df[input].values.tolist(), model_id, + n=n, **_clean_kwargs(kwargs) ) - for out in output: - df[out] = data + if n and n > 1 and len(output) == n: + # Distribute: each output column gets the nth match + _distribute_n_matches(data) + elif n and n > 1 and len(output) > 1: + raise ValueError( + f'When n > 1 and multiple output columns are provided, the number ' + f'of output columns ({len(output)}) must equal n ({n}).' + ) + else: + for out in output: + df[out] = data else: - # User specified a mixture of unrecognized columns and columns from the wrangle + # User specified a mixture of unrecognized columns and columns from the wrangle raise ValueError('Lookup may only contain all named or unnamed columns.') elif lookup_mode == 'by_dataframe': @@ -1593,6 +1640,39 @@ def resolve_input(candidate_list): ): del kwargs["functions"] + def output_exists(output_column): + return output_column in df.columns + + def resolve_rename_input(input_column, output_column): + candidates = input_column if isinstance(input_column, list) else [input_column] + optional_candidates = [] + required_candidates = [] + + for candidate in candidates: + optional = False + actual_col = candidate + if isinstance(candidate, str) and candidate.endswith("?"): + optional = True + actual_col = candidate[:-1] + + if actual_col in df.columns: + return actual_col + + if optional: + optional_candidates.append(actual_col) + continue + + required_candidates.append(actual_col) + + if optional_candidates and not required_candidates: + return None + + if output_exists(output_column): + return None + + missing = required_candidates[0] if required_candidates else candidates[0] + raise ValueError(f'Rename column "{missing}" not found.') + # If short form of paired names is provided, use that if input is None: @@ -1650,7 +1730,7 @@ def resolve_input(candidate_list): # Regular non-wildcard name if name not in cols: - if optional: + if optional or kwargs[x] in cols: continue else: raise ValueError(f'Rename column "{name}" not found.') @@ -1676,19 +1756,12 @@ def resolve_input(candidate_list): raise ValueError('The lists for input and output must be the same length.') for inp, out in zip(input, output): - if inp.endswith("?"): - actual_col = inp[:-1] - if actual_col not in list(df.columns): - # Skip this column if it doesn't exist - continue # This skips both input and output - else: - filtered_input.append(actual_col) - filtered_output.append(out) - elif inp not in list(df.columns): - raise ValueError(f'Rename column "{inp}" not found.') - else: - filtered_input.append(inp) - filtered_output.append(out) + actual_col = resolve_rename_input(inp, out) + if actual_col is None: + continue + + filtered_input.append(actual_col) + filtered_output.append(out) # Check that the output columns don't already exist if so drop them df = df.drop(columns=[x for x in filtered_output if x in df.columns and x not in filtered_input]) diff --git a/wrangles/recipe_wrangles/merge.py b/wrangles/recipe_wrangles/merge.py index c79d42b36..cc8d7627d 100644 --- a/wrangles/recipe_wrangles/merge.py +++ b/wrangles/recipe_wrangles/merge.py @@ -3,9 +3,11 @@ """ from typing import Union as _Union import fnmatch as _fnmatch -import numpy as _np + +import logging as _logging import pandas as _pd from .. import format as _format +import numpy as _np def coalesce( @@ -36,6 +38,7 @@ def coalesce( # Ensure input is a list if not isinstance(input, list): input = [input] + _logging.debug(f": Coalescing values :: input :: {input}") if len(input) == 1: if output is None: output = input[0] @@ -54,7 +57,7 @@ def coalesce( df[output] = _pd.Series(dtype=object) return df - arr = df[input].fillna('').values # object array (n_rows, n_cols) + arr = df[input].fillna('').to_numpy(dtype=object) # object array (n_rows, n_cols) n_rows = len(arr) m = len(input) @@ -78,7 +81,7 @@ def _safe_strip(x): try: return x.strip() except AttributeError: - return '' if not x else x + return x # non-strings can't be whitespace-only stripped_result = _np.frompyfunc(_safe_strip, 1, 1)(result) needs_fallback = _np.where(has_any & (stripped_result == ''))[0] for i in needs_fallback: @@ -86,7 +89,10 @@ def _safe_strip(x): val = arr[i, j] if isinstance(val, str): val = val.strip() - if val: + if val: + result[i] = val + break + else: result[i] = val break else: @@ -136,6 +142,7 @@ def concatenate( if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] + _logging.debug(f": Concatenating columns :: input :: {input}, char :: {char}") if len(input) == 1: df[output[0]] = _format.concatenate(df[input[0]].values, char, skip_empty) else: @@ -165,6 +172,7 @@ def dictionaries(df: _pd.DataFrame, input: list, output: str, skip_empty: bool = default: false """ + _logging.debug(f": Merging dictionary columns :: input :: {input}") rows = df[input].values.tolist() if skip_empty: @@ -208,6 +216,7 @@ def key_value_pairs(df: _pd.DataFrame, input: dict, output: str, skip_empty: boo description: Whether to skip empty keys or values when creating the dictionary default: false """ + _logging.debug(f": Creating key-value pairs :: input :: {input}") pairs = {} # If user has used wildcards, expand out @@ -287,6 +296,7 @@ def lists(df: _pd.DataFrame, input: list, output: str, remove_duplicates: bool = type: boolean description: Whether to include empty values in the created list """ + _logging.debug(f": Merging list columns :: input :: {input}") output_list = [] for row in df[input].values.tolist(): output_row = [] @@ -334,6 +344,7 @@ def to_dict(df: _pd.DataFrame, input: list, output: str, include_empty: bool = F type: boolean description: Whether to include empty columns in the created dictionary """ + _logging.debug(f": Converting columns to dict :: output :: {output}") index_check = 0 cols_changed = [] for cols in input: @@ -386,6 +397,7 @@ def to_list(df: _pd.DataFrame, input: list, output: str, include_empty: bool = F type: boolean description: Whether to include empty columns in the created list """ + _logging.debug(f": Converting columns to list :: output :: {output}") output_list = [] for row in df[input].values.tolist(): output_row = [] diff --git a/wrangles/recipe_wrangles/pandas.py b/wrangles/recipe_wrangles/pandas.py index 3078c6880..aadd33d31 100644 --- a/wrangles/recipe_wrangles/pandas.py +++ b/wrangles/recipe_wrangles/pandas.py @@ -1,6 +1,57 @@ import pandas as _pd from typing import Union as _Union from numpy import nan as _nan +import logging as _logging + + +def _sort_by_columns(by): + return [by] if isinstance(by, str) else list(by) + + +def _is_empty_sort_value(value): + if isinstance(value, str) and value.strip() == "": + return True + + try: + return bool(_pd.isna(value)) + except (TypeError, ValueError): + return False + + +def _coerce_sort_series(series: _pd.Series) -> _pd.Series: + non_empty = series[~series.map(_is_empty_sort_value)] + if non_empty.empty: + return series + + numeric_values = _pd.to_numeric(series, errors="coerce") + numeric_count = numeric_values.loc[non_empty.index].notna().sum() + if numeric_count >= len(non_empty) / 2: + return numeric_values.fillna(0) + + return series.map(lambda value: "" if _is_empty_sort_value(value) else str(value)) + + +def _coerced_sort_values(df: _pd.DataFrame, by_columns: list, ignore_index: bool, kwargs: dict) -> _pd.DataFrame: + sort_df = df.copy() + sort_kwargs = kwargs.copy() + temp_columns = [] + + for index, column in enumerate(by_columns): + if column not in sort_df.columns: + continue + + temp_column = f"__wrangles_sort_key_{index}" + while temp_column in sort_df.columns: + temp_column = f"_{temp_column}" + + sort_df[temp_column] = _coerce_sort_series(sort_df[column]) + temp_columns.append(temp_column) + + if not temp_columns: + return df.sort_values(ignore_index=ignore_index, **kwargs) + + sort_kwargs["by"] = temp_columns + return sort_df.sort_values(ignore_index=ignore_index, **sort_kwargs).drop(columns=temp_columns) def copy( @@ -29,6 +80,7 @@ def copy( - array description: Name of the output columns or columns """ + _logging.debug(f": Copying columns :: input :: {input}") # If short form of paired names is provided, use that if input is None: # Check that column name exists @@ -75,6 +127,7 @@ def drop(df: _pd.DataFrame, columns: _Union[str, list]) -> _pd.DataFrame: - string description: Name of the column(s) to drop """ + _logging.debug(f": Dropping columns :: {columns}") return df.drop(columns=columns, errors='ignore') @@ -94,6 +147,7 @@ def transpose(df: _pd.DataFrame, header_column = 0) -> _pd.DataFrame: for the transposed DataFrame. Default 0 (first column). Use header_column = null to not use any column as header. """ + _logging.debug(f": Transposing dataframe :: header_column :: {header_column}") if header_column is not None: if isinstance(header_column, int): # If header_column is an integer, use it as the column index @@ -137,13 +191,14 @@ def sort(df: _pd.DataFrame, ignore_index=True, **kwargs) -> _pd.DataFrame: If this is a list of bools then it must match the length of the by. """ + _logging.debug(": Sorting dataframe") # Extract and normalize the 'by' parameter by = kwargs.get("by") if by is None: raise ValueError("'by' parameter is required for sorting") # Ensure 'by' is a list for consistent processing - by_columns = [by] if isinstance(by, str) else by + by_columns = _sort_by_columns(by) # Check if any columns need dtype conversion (float16 -> float32) # float16 can cause sorting issues in pandas @@ -161,8 +216,12 @@ def sort(df: _pd.DataFrame, ignore_index=True, **kwargs) -> _pd.DataFrame: if col in df.columns and df[col].dtype == "float16": df[col] = df[col].astype("float32") - # Perform the sort operation - return df.sort_values(ignore_index=ignore_index, **kwargs) + # Perform the sort operation. If mixed column types cannot be compared, + # retry using temporary sort keys coerced to the predominant compatible type. + try: + return df.sort_values(ignore_index=ignore_index, **kwargs) + except TypeError: + return _coerced_sort_values(df, by_columns, ignore_index, kwargs) def round(df: _pd.DataFrame, input: _Union[str, int, list], decimals: int = 0, output: _Union[str, list] = None) -> _pd.DataFrame: @@ -194,6 +253,7 @@ def round(df: _pd.DataFrame, input: _Union[str, int, list], decimals: int = 0, o if not isinstance(input, list): input = [input] if not isinstance(output, list): output = [output] + _logging.debug(f": Rounding columns :: input :: {input}, decimals :: {decimals}") for input_column, output_column in zip(input, output): # coerce input column to floats (nan on error) # replace nan with empty string @@ -277,6 +337,7 @@ def explode( If false, rows that contain empty lists will keep 1 row with an empty value. Default False. """ + _logging.debug(f": Exploding columns :: {input}") # If a string provided, convert to list if not isinstance(input, list): input = [input] @@ -295,4 +356,4 @@ def explode( if drop_empty: df = df.dropna(subset=input, how='all') - return df \ No newline at end of file + return df diff --git a/wrangles/recipe_wrangles/select.py b/wrangles/recipe_wrangles/select.py index 5b8d53a5d..d170e83ca 100644 --- a/wrangles/recipe_wrangles/select.py +++ b/wrangles/recipe_wrangles/select.py @@ -5,6 +5,7 @@ import types as _types import re as _re import json as _json +import logging as _logging import pandas as _pd from .. import select as _select from ..utils import get_nested_function as _get_nested_function @@ -24,11 +25,12 @@ def columns(df: _pd.DataFrame, input: _Union[str, int, list]) -> _pd.DataFrame: - array description: Name of the column(s) to select """ + _logging.debug(f": Selecting columns :: input :: {input}") if not isinstance(input, list): input = [input] - + # Missing column should be caught by _wildcard_expansion - - return df[input] + + return df[input] def dictionary_element( @@ -90,6 +92,7 @@ def dictionary_element( if len(input) != len(output): raise ValueError('The list of inputs and outputs must be the same length for select.dictionary_element') + _logging.debug(f": Selecting dictionary element :: {element} from {input}") for in_col, out_col in zip(input, output): df[out_col] = _select.dict_element(df[in_col].tolist(), element, default=default) @@ -168,6 +171,8 @@ def _extract_elements(input_string): # Ensure input and output are equal lengths if len(input) != len(output): raise ValueError('The list of inputs and outputs must be the same length for select.element') + + _logging.debug(f": Selecting elements :: input :: {input}") # Handle default values - convert to list if needed if default is None: @@ -507,6 +512,7 @@ def head(df: _pd.DataFrame, n: int) -> _pd.DataFrame: type: integer description: Number of rows to return """ + _logging.debug(f": Selecting head :: n :: {n}") if not isinstance(n, int) or n <= 0: raise ValueError("n must be a positive integer") return df.head(n) @@ -531,6 +537,7 @@ def highest_confidence(df: _pd.DataFrame, input: list, output: _Union[str,list]) description: If two columns; the result and confidence. If one column; [result, confidence] """ + _logging.debug(f": Selecting highest confidence :: input :: {input}") if isinstance(output, list) and len(output) == 1: output = output[0] elif isinstance(output, list) and len(output) > 2: @@ -704,6 +711,7 @@ def list_element( if len(input) != len(output): raise ValueError('The list of inputs and outputs must be the same length for select.list_element') + _logging.debug(f": Selecting list element :: {element} from {input}") for in_col, out_col in zip(input, output): df[out_col] = _select.list_element(df[in_col].tolist(), element, default=default) @@ -933,5 +941,6 @@ def threshold(df: _pd.DataFrame, input: list, output: str, threshold: float) -> minimum: 0 maximum: 1 """ + _logging.debug(f": Applying confidence threshold :: {threshold} on {input}") df[output] = _select.confidence_threshold(df[input[0]].tolist(), df[input[1]].tolist(), threshold) return df diff --git a/wrangles/recipe_wrangles/split.py b/wrangles/recipe_wrangles/split.py index dee020e10..01f18cd06 100644 --- a/wrangles/recipe_wrangles/split.py +++ b/wrangles/recipe_wrangles/split.py @@ -3,6 +3,7 @@ """ # Rename List to _list to be able to use function name list without clashing from typing import Union as _Union, List as _list +import logging as _logging import pandas as _pd from .. import format as _format import json as _json @@ -16,7 +17,8 @@ def dictionary( df: _pd.DataFrame, input: _Union[str, int, _list], output: _Union[str, _list] = None, - default: dict = None + default: dict = None, + output_format: str = "columns" ) -> _pd.DataFrame: """ type: object @@ -29,7 +31,7 @@ def dictionary( - input properties: input: - type: + type: - string - integer - array @@ -38,24 +40,39 @@ def dictionary( If providing multiple dictionaries and the dictionaries contain overlapping values, the last value will be returned. output: - type: + type: - string - array description: |- - (Optional) Subset of keys to extract from the dictionary. - If not provided, all keys will be returned. + In columns output_format, this is an optional subset of keys to extract + from the dictionary. If not provided, all keys will be returned. Columns can be renamed with the following syntax: output: - key1: new_column_name1 - key2: new_column_name2 + In to_lists output_format, this must be two output columns for the keys + and values lists. If not provided, Keys and Values will be used. default: type: object description: >- Provide a set of default headings and values if they are not found within the input - """ + output_format: + type: string + enum: + - columns + - to_lists + description: |- + How to split the dictionary. + columns creates one output column for each dictionary key. + to_lists creates two output columns containing lists of keys and values. + """ + if output_format not in ["columns", "to_lists"]: + raise ValueError("output_format must be one of: columns, to_lists") + if default is None: default = {} + _logging.debug(f": Splitting dictionaries :: input :: {input}") # Ensure input is passed as a list if not isinstance(input, _list): input = [input] @@ -71,11 +88,28 @@ def _parse_dict_or_json(val): raise ValueError(f'{val} is not a valid Dictionary') from None - # Generate new columns for each key in the dictionary - df_temp = _pd.DataFrame([ + # Merge each row's dictionaries so duplicate keys follow existing behavior: + # later input columns overwrite earlier input columns. + dicts = [ dict(_itertools.chain.from_iterable(_parse_dict_or_json(d) for d in ([default] + row.tolist()))) for row in df[input].values - ]) + ] + + if output_format == "to_lists": + if output is None: + output = ["Keys", "Values"] + elif not isinstance(output, _list): + output = [output] + + if len(output) != 2: + raise ValueError("split.dictionary with output_format to_lists requires exactly two output columns") + + df[output[0]] = [[key for key in item.keys()] for item in dicts] + df[output[1]] = [[value for value in item.values()] for item in dicts] + return df + + # Generate new columns for each key in the dictionary + df_temp = _pd.DataFrame(dicts) # If user has defined how they'd like the output columns if output is not None: @@ -130,6 +164,7 @@ def list(df: _pd.DataFrame, input: _Union[str, int], output: _Union[str, _list]) If providing a single column, use a wildcard (*) to indicate a incrementing integer """ + _logging.debug(f": Splitting lists :: input :: {input}") # Ensure rows are lists even if they are JSON strings results = [ row if isinstance(row, _list) else _json.loads(row) @@ -218,6 +253,7 @@ def text( description: Whether to skip empty values default: false """ + _logging.debug(f": Splitting text :: char :: {char}, input :: {input}") # Ensure only a single input column is specified if isinstance(input, _list): if len(input) != 1: @@ -323,6 +359,8 @@ def tokenize( # Ensure input and output are equal lengths if len(input) != len(output): raise ValueError('The list of inputs and outputs must be the same length for split.tokenize') + + _logging.debug(f": Tokenizing :: input :: {input}") func = None pattern = None diff --git a/wrangles/select.py b/wrangles/select.py index e1bd83baa..725ca1fe3 100644 --- a/wrangles/select.py +++ b/wrangles/select.py @@ -5,6 +5,7 @@ import json as _json import itertools as _itertools import fnmatch as _fnmatch +import logging as _logging from .utils import wildcard_expansion_dict as _wildcard_expansion_dict from .utils import wildcard_expansion as _wildcard_expansion @@ -12,6 +13,7 @@ def highest_confidence(data_list): """ Select the option with the highest confidence from multiple columns """ + _logging.debug(f": Selecting highest confidence from {len(data_list)} records") results = [] for row in data_list: highest_confidence = 0 @@ -53,6 +55,7 @@ def confidence_threshold(list_1, list_2, threshold): """ Select the first option if it exceeds a given threshold, else the second option. """ + _logging.debug(f": Applying confidence threshold :: {threshold}") results = [] for cell_1, cell_2 in zip(list_1, list_2): @@ -78,6 +81,7 @@ def list_element(input, n: _Union[str, int], default = ""): """ Select a numbered element of a list (zero indexed). """ + _logging.debug(f": Extracting element {n} from {len(input)} lists") def _int_or_none(val): try: return int(val) @@ -111,6 +115,7 @@ def dict_element(input: _Union[list, dict], key: _Union[str, list], default: any """ Select an element or elements of a dictionary """ + _logging.debug(f": Extracting dict element :: {key} from {len(input) if isinstance(input, list) else 1} records") # Ensure input is a list single_input = False if not isinstance(input, list): diff --git a/wrangles/standardize.py b/wrangles/standardize.py index 474a2496c..37447c578 100644 --- a/wrangles/standardize.py +++ b/wrangles/standardize.py @@ -1,4 +1,5 @@ from typing import Union as _Union +import logging as _logging from . import config as _config from . import data as _data from . import batching as _batching @@ -52,6 +53,7 @@ def standardize( if purpose != 'standardize': raise ValueError(f'Using {purpose} model_id {model_id} in a standardize function.') + _logging.info(f": Standardizing {len(json_data)} records :: model_id :: {model_id}, case_sensitive :: {case_sensitive}") results = _batching.batch_api_calls(url, params, json_data, batch_size) if isinstance(input, str): results = results[0] diff --git a/wrangles/train.py b/wrangles/train.py index 9455fad35..9456d7663 100644 --- a/wrangles/train.py +++ b/wrangles/train.py @@ -261,6 +261,23 @@ def lookup( return response + def delete(model_id: str): + """ + Delete a trained model by its model ID. + Requires WrangleWorks Account and Subscription. + + :param model_id: The ID of the model to delete. + """ + _logging.info(f": Deleting model :: {model_id}") + response = _requests.delete( + f'{_config.api_host}/model/delete', + params={'model_id': model_id}, + headers={'Authorization': f'Bearer {_auth.get_access_token()}'} + ) + if not response.ok: + raise RuntimeError(f"Delete model failed. {response.status_code} : {response.text}") + return response + def standardize(training_data: list, name: str = None, model_id: str = None): """ Train a standardize model. This can standardize text to a desired format. diff --git a/wrangles/translate.py b/wrangles/translate.py index e67b1e538..72269beac 100644 --- a/wrangles/translate.py +++ b/wrangles/translate.py @@ -2,6 +2,7 @@ Functions to translate text """ from typing import Union as _Union +import logging as _logging from . import config as _config from . import batching as _batching @@ -102,6 +103,7 @@ def translate( elif case == 'title': json_data = [val.title() for val in json_data] + _logging.info(f": Translating {len(json_data)} records :: {source_language} -> {target_language}") url = f'{_config.api_host}/wrangles/translate' params = { 'responseFormat':'array', diff --git a/wrangles/utils.py b/wrangles/utils.py index 334324c1d..4fc84d80b 100644 --- a/wrangles/utils.py +++ b/wrangles/utils.py @@ -254,6 +254,8 @@ def escape_except(text, chars_not_to_escape): if not isinstance(selected_columns, list): selected_columns = [selected_columns] + _logging.debug(f": Expanding wildcards :: {len(selected_columns)} patterns against {len(all_columns)} columns") + # Convert wildcards to regex pattern for i, val in enumerate(selected_columns): if val in all_columns: @@ -445,6 +447,7 @@ def request_retries(request_type, url, **kwargs): :param kwargs: Arguments to pass to requests.request :returns: requests.Response object """ + _logging.debug(f": HTTP request :: method :: {request_type}, url :: {url}") session = _requests.Session() session.mount( 'https://', @@ -655,6 +658,7 @@ def __getattr__(self, item): except ImportError as e: raise ImportError( f"Optional dependency '{self.module_name}' is required for this feature. " - f"Please install it with: pip install {self.module_name}" + f"Please install it with: pip install {self.module_name} " + f"or install all optional dependencies with: pip install -r requirements-full.txt" ) from e return getattr(self._module, item)