(feat, python): support optional python deps + extras (#3742)

generators/python/sdk/CHANGELOG.md

@@ -5,6 +5,21 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.6.0] - 2024-05-30
+
+- Improvement: Support adding optional dependencies and extras to your generated `pyproject.toml`. To 
+  use this configuration, please add the following: 
+
+  ```yaml
+  extra_dependencies:
+    boto3: 1.28.57
+    langchain: 
+      version: "^0.1.20"
+      optional: true
+  extras: 
+    telemetry: ["langchain", "boto3"]
+  ```
+
 ## [2.5.7] - 2024-05-30
 
 - Fix: tests now carry a type annotation for `expected_types` variable.

generators/python/src/fern_python/codegen/ast/dependency/dependency.py

@@ -15,3 +15,4 @@ class Dependency:
     name: DependencyName
     version: DependencyVersion
     compatibility: DependencyCompatibility = DependencyCompatibility.EXACT
+    optional: bool = False

generators/python/src/fern_python/codegen/project.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import typing
 from dataclasses import dataclass
 from pathlib import Path
 from types import TracebackType
@@ -67,6 +68,7 @@ def __init__(
         self._github_output_mode = github_output_mode
         self._pypi_metadata = pypi_metadata
         self.license_ = license_
+        self._extras: typing.Dict[str, List[str]] = {}
 
     def add_init_exports(self, path: AST.ModulePath, exports: List[ModuleExport]) -> None:
         self._module_manager.register_additional_exports(path, exports)
@@ -77,6 +79,9 @@ def add_dependency(self, dependency: AST.Dependency) -> None:
     def add_dev_dependency(self, dependency: AST.Dependency) -> None:
         self._dependency_manager.add_dev_dependency(dependency)
 
+    def add_extra(self, extra: typing.Dict[str, List[str]]) -> None:
+        self._extras = extra
+
     def set_generate_readme(self, generate_readme: bool) -> None:
         self._generate_readme = generate_readme
 
@@ -167,6 +172,7 @@ def finish(self) -> None:
                 github_output_mode=self._github_output_mode,
                 pypi_metadata=self._pypi_metadata,
                 license_=self.license_,
+                extras=self._extras,
             )
             py_project_toml.write()
 

generators/python/src/fern_python/codegen/pyproject_toml.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+import typing
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import List, Optional, Set, cast
@@ -40,6 +41,7 @@ def __init__(
         pypi_metadata: Optional[PypiMetadata],
         github_output_mode: Optional[GithubOutputMode],
         license_: Optional[LicenseConfig],
+        extras: typing.Dict[str, List[str]] = {},
     ):
         self._poetry_block = PyProjectToml.PoetryBlock(
             name=name,
@@ -52,6 +54,7 @@ def __init__(
         self._dependency_manager = dependency_manager
         self._path = path
         self._python_version = python_version
+        self._extras = extras
 
     def write(self) -> None:
         blocks: List[PyProjectToml.Block] = [
@@ -67,6 +70,15 @@ def write(self) -> None:
         content = ""
         for block in blocks:
             content += block.to_string()
+
+        if len(self._extras) > 0:
+            content += f"""
+[tool.poetry.extras]
+"""
+            for key, vals in self._extras.items():
+                stringified_vals = ", ".join([f'"{val}"' for val in vals])
+                content += f"{key}=[{stringified_vals}]\n"
+
         with open(os.path.join(self._path, "pyproject.toml"), "w") as f:
             f.write(content)
 
@@ -177,14 +189,16 @@ def deps_to_string(self, dependencies: Set[Dependency]) -> str:
             deps = ""
             for dep in sorted(dependencies, key=lambda dep: dep.name):
                 compatiblity = dep.compatibility
-                # TODO(dsinghvi): assert all enum cases are visited
-                print(dep.compatibility)
-                if compatiblity == DependencyCompatibility.EXACT:
-                    print(f"{dep.name} is exact")
-                    deps += f'{dep.name.replace(".", "-")} = "{dep.version}"\n'
-                elif compatiblity == DependencyCompatibility.GREATER_THAN_OR_EQUAL:
-                    print(f"{dep.name} is greater than or equal")
-                    deps += f'{dep.name.replace(".", "-")} = ">={dep.version}"\n'
+                is_optional = dep.optional
+                version = dep.version
+                name = dep.name.replace(".", "-")
+                if compatiblity == DependencyCompatibility.GREATER_THAN_OR_EQUAL:
+                    version = f">={dep.version}"
+
+                if is_optional:
+                    deps += f'{name} = {{ version="{version}", optional = true}}\n'
+                else:
+                    deps += f'{name} = "{version}"\n'
             return deps
 
         def to_string(self) -> str:

generators/python/src/fern_python/generators/sdk/custom_config.py

@@ -28,9 +28,15 @@ class Config:
         extra = pydantic.Extra.forbid
 
 
+class DependencyCusomConfig(pydantic.BaseModel):
+    version: str
+    optional: bool
+
+
 class SDKCustomConfig(pydantic.BaseModel):
-    extra_dependencies: Dict[str, str] = {}
+    extra_dependencies: Dict[str, Union[str, DependencyCusomConfig]] = {}
     extra_dev_dependencies: Dict[str, str] = {}
+    extras: Dict[str, List[str]] = {}
     skip_formatting: bool = False
     client: ClientConfiguration = ClientConfiguration()
     include_union_utils: bool = False

generators/python/src/fern_python/generators/sdk/sdk_generator.py

@@ -29,7 +29,7 @@
 from .client_generator.generated_root_client import GeneratedRootClient
 from .client_generator.oauth_token_provider_generator import OAuthTokenProviderGenerator
 from .client_generator.root_client_generator import RootClientGenerator
-from .custom_config import SDKCustomConfig
+from .custom_config import DependencyCusomConfig, SDKCustomConfig
 from .environment_generators import (
     GeneratedEnvironment,
     MultipleBaseUrlsEnvironmentGenerator,
@@ -84,8 +84,15 @@ def run(
         if not custom_config.client.exported_filename.endswith(".py"):
             raise RuntimeError("client_location.exported_filename must end in .py")
 
-        for dep, version in custom_config.extra_dependencies.items():
-            project.add_dependency(dependency=AST.Dependency(name=dep, version=version))
+        for dep, value in custom_config.extra_dependencies.items():
+            if type(value) is str:
+                project.add_dependency(dependency=AST.Dependency(name=dep, version=value))
+            elif isinstance(value, DependencyCusomConfig):
+                project.add_dependency(
+                    dependency=AST.Dependency(name=dep, version=value.version, optional=value.optional)
+                )
+
+        project.add_extra(custom_config.extras)
 
         for dep, version in custom_config.extra_dev_dependencies.items():
             project.add_dev_dependency(dependency=AST.Dependency(name=dep, version=version))

packages/seed/src/cli.ts

@@ -51,6 +51,11 @@ function addTestCommand(cli: Argv) {
                     demandOption: false,
                     description: "Runs on all fixtures if not provided"
                 })
+                .option("outputFolder", {
+                    string: true,
+                    demandOption: false,
+                    description: "Runs on a specific output folder. Only relevant if there are >1 folders configured."
+                })
                 .option("keepDocker", {
                     type: "boolean",
                     demandOption: false,
@@ -112,7 +117,8 @@ function addTestCommand(cli: Argv) {
                     testGenerator({
                         generator,
                         runner: testRunner,
-                        fixtures: argv.fixture
+                        fixtures: argv.fixture,
+                        outputFolder: argv.outputFolder
                     })
                 );
             }

packages/seed/src/commands/test/testWorkspaceFixtures.ts

@@ -16,17 +16,22 @@ export const FIXTURES = readDirectories(path.join(__dirname, FERN_DIRECTORY, API
 export async function testGenerator({
     runner,
     generator,
-    fixtures
+    fixtures,
+    outputFolder
 }: {
     runner: TestRunner;
     generator: GeneratorWorkspace;
     fixtures: string[];
+    outputFolder?: string;
 }): Promise<boolean> {
     const testCases: Promise<TestRunner.TestResult>[] = [];
     for (const fixture of fixtures) {
         const config = generator.workspaceConfig.fixtures?.[fixture];
         if (config != null) {
             for (const instance of config) {
+                if (outputFolder != null && instance.outputFolder !== outputFolder) {
+                    continue;
+                }
                 testCases.push(
                     runner.run({
                         fixture,

seed/python-sdk/seed.yml

@@ -44,8 +44,14 @@ fixtures:
           extra_fields: "allow"
       outputFolder: pydantic-extra-fields
     - customConfig:
+        pyproject_python_version: "^3.8.1"
         extra_dependencies:
           boto3: 1.28.57
+          langchain: 
+            version: "^0.1.20"
+            optional: true
+        extras: 
+          telemetry: ["langchain", "boto3"]
       outputFolder: extra_dependencies
     - customConfig:
         pyproject_python_version: "^3.8.1"