This file is a merged representation of the entire codebase, combined into a single document by Repomix.
The content has been processed where content has been compressed (code blocks are separated by ⋮---- delimiter).
This section contains a summary of this file.
This file contains a packed representation of the entire repository's contents.
It is designed to be easily consumable by AI systems for analysis, code review,
or other automated processes.
The content is organized as follows:
1. This summary section
2. Repository information
3. Directory structure
4. Repository files (if enabled)
5. Multiple file entries, each consisting of:
- File path as an attribute
- Full contents of the file
- This file should be treated as read-only. Any changes should be made to the
original repository files, not this packed version.
- When processing this file, use the file path to distinguish
between different files in the repository.
- Be aware that this file may contain sensitive information. Handle it with
the same level of security as you would the original repository.
- Some files may have been excluded based on .gitignore rules and Repomix's configuration
- Binary files are not included in this packed representation. Please refer to the Repository Structure section for a complete list of file paths, including binary files
- Files matching patterns in .gitignore are excluded
- Files matching default ignore patterns are excluded
- Content has been compressed - code blocks are separated by ⋮---- delimiter
- Files are sorted by Git change count (files with more changes are at the bottom)
mix/
tasks/
compile/
snakebridge.ex
snakebridge.setup.ex
snakebridge.verify.ex
snakebridge/
docs/
markdown_converter.ex
math_renderer.ex
rst_parser.ex
error/
dtype_mismatch_error.ex
out_of_memory_error.ex
shape_mismatch_error.ex
generator/
type_mapper.ex
lock/
verifier.ex
python_runner/
system.ex
telemetry/
handlers/
logger.ex
metrics.ex
runtime_forwarder.ex
types/
decoder.ex
encoder.ex
wheel_selector/
config_strategy.ex
adapter.ex
application.ex
benchmark.ex
bytes.ex
callback_registry.ex
compile_error.ex
config.ex
defaults.ex
docs.ex
dynamic_exception.ex
dynamic.ex
environment_error.ex
error_translator.ex
error.ex
examples.ex
generator.ex
helper_generator.ex
helper_not_found_error.ex
helper_registry_error.ex
helpers.ex
introspection_error.ex
introspector.ex
invalid_ref_error.ex
ledger.ex
lock.ex
manifest.ex
module_resolver.ex
python_env.ex
python_packages_runner.ex
python_runner.ex
python_runtime_runner.ex
ref_not_found_error.ex
ref.ex
registry.ex
runtime_client.ex
runtime.ex
scan_error.ex
scanner.ex
serialization_error.ex
session_context.ex
session_manager.ex
session_mismatch_error.ex
snakepit_types.ex
stream_ref.ex
telemetry.ex
types.ex
wheel_config.ex
wheel_selector.ex
with_context.ex
snakebridge.ex
This section contains the contents of the repository's files.
defmodule Mix.Tasks.Compile.Snakebridge do
@moduledoc """
Mix compiler that runs the SnakeBridge pre-pass (scan → introspect → generate).
"""
use Mix.Task.Compiler
alias SnakeBridge.{
Config,
Generator,
HelperGenerator,
Helpers,
Introspector,
Lock,
Manifest,
ModuleResolver,
PythonEnv,
Scanner,
Telemetry
}
@reserved_words ~w(def defp defmodule class do end if unless case cond for while with fn when and or not true false nil in try catch rescue after else raise throw receive)
@impl Mix.Task.Compiler
def run(_args) do
Mix.Task.run("loadconfig")
if skip_generation?() do
{:ok, []}
else
config = Config.load()
run_with_config(config)
end
end
defp run_with_config(%{libraries: []}), do: {:ok, []}
defp run_with_config(config) do
if strict_mode?(config) do
run_strict(config)
else
PythonEnv.ensure!(config)
run_normal(config)
end
end
@impl Mix.Task.Compiler
def manifests do
Mix.Task.run("loadconfig")
config = Config.load()
[
Path.join(config.metadata_dir, "manifest.json"),
"snakebridge.lock"
]
end
defp skip_generation? do
case System.get_env("SNAKEBRIDGE_SKIP") do
nil -> false
value -> value in ["1", "true", "TRUE", "yes", "YES"]
end
end
defp update_manifest(manifest, targets) do
{updated_manifest, errors} =
targets
|> Introspector.introspect_batch()
|> Enum.reduce({manifest, []}, fn {library, result, python_module}, {acc, errs} ->
case result do
{:ok, infos} ->
{symbol_entries, class_entries} =
build_manifest_entries(library, python_module, infos)
updated =
acc
|> Manifest.put_symbols(symbol_entries)
|> Manifest.put_classes(class_entries)
{updated, errs}
{:error, reason} ->
log_introspection_error(library, python_module, reason)
emit_introspection_error_telemetry(library, python_module, reason)
{acc, [{library, python_module, reason} | errs]}
end
end)
if errors != [] do
show_introspection_summary(errors)
end
updated_manifest
end
defp log_introspection_error(library, python_module, reason) do
formatted = format_introspection_error(library, python_module, reason)
Mix.shell().info(formatted)
end
defp format_introspection_error(library, python_module, reason) do
library_name = get_library_name(library)
base = build_base_message(library_name, python_module)
format_reason(base, reason)
end
defp get_library_name(library) when is_map(library), do: library.name || library.python_name
defp get_library_name(library), do: inspect(library)
defp build_base_message(library_name, python_module) do
base = " [warning] Introspection failed for #{library_name}"
if python_module && python_module != library_name do
base <> ".#{python_module}"
else
base
end
end
defp format_reason(base, %{type: _type, message: message, suggestion: suggestion}) do
lines = [base, " Error: #{message}"]
lines = if suggestion, do: lines ++ [" Suggestion: #{suggestion}"], else: lines
Enum.join(lines, "\n")
end
defp format_reason(base, %{message: message}) do
base <> "\n Error: #{message}"
end
defp format_reason(base, message) when is_binary(message) do
base <> "\n Error: #{message}"
end
defp format_reason(base, reason) do
base <> "\n Error: #{inspect(reason)}"
end
defp emit_introspection_error_telemetry(library, python_module, reason) do
library_name =
if is_map(library), do: library.name || library.python_name, else: inspect(library)
error_type =
case reason do
%{type: type} -> type
_ -> :unknown
end
:telemetry.execute(
[:snakebridge, :introspection, :error],
%{count: 1},
%{
library: library_name,
python_module: python_module,
error_type: error_type,
reason: reason
}
)
end
defp show_introspection_summary(errors) do
count = length(errors)
message = """
================================================================================
SnakeBridge Introspection Summary
================================================================================
#{count} introspection error(s) occurred. Some symbols may be missing from
the generated bindings.
To resolve:
1. Check the errors above for details
2. Ensure Python packages are installed: mix snakebridge.setup
3. Check for import errors in your Python dependencies
4. Re-run: mix compile
The compilation will continue, but affected symbols will not be available.
================================================================================
"""
Mix.shell().info(message)
end
defp build_targets(missing, config, manifest) do
initial =
Enum.reduce(missing, %{}, fn entry, acc ->
accumulate_missing_target(entry, acc, config)
end)
with_includes =
Enum.reduce(config.libraries, initial, fn library, acc ->
includes =
library.include
|> Enum.reject(&function_or_class_present?(manifest, library, &1))
accumulate_includes(includes, library, acc)
end)
with_includes
|> Enum.map(fn {{library, python_module}, functions} ->
filtered = Enum.reject(functions, &(&1 in library.exclude))
{library, python_module, Enum.uniq(filtered)}
end)
|> Enum.reject(fn {_library, _python_module, functions} -> functions == [] end)
end
defp accumulate_missing_target({module, function, _arity}, acc, config) do
case library_for_module(config, module) do
nil ->
acc
library ->
case ModuleResolver.resolve_class_or_submodule(library, module) do
{:class, class_name, parent_module} ->
add_target(acc, library, parent_module, class_name)
{:submodule, python_module} ->
python_function =
function
|> to_string()
|> python_name_for_elixir_name()
add_target(acc, library, python_module, python_function)
{:error, _reason} ->
python_module = python_module_for_elixir(library, module)
python_function =
function
|> to_string()
|> python_name_for_elixir_name()
add_target(acc, library, python_module, python_function)
end
end
end
defp accumulate_includes([], _library, acc), do: acc
defp accumulate_includes(includes, library, acc) do
key = {library, library.python_name}
Map.update(acc, key, includes, fn funcs -> includes ++ funcs end)
end
defp add_target(acc, library, python_module, function) do
key = {library, python_module}
Map.update(acc, key, [function], fn funcs -> [function | funcs] end)
end
defp library_for_module(config, module) do
module_parts = Module.split(module)
Enum.find(config.libraries, fn library ->
library_parts = Module.split(library.module_name)
Enum.take(module_parts, length(library_parts)) == library_parts
end)
end
defp build_manifest_entries(library, python_module, infos) do
Enum.reduce(infos, {[], []}, fn info, {symbols, classes} ->
cond do
info["error"] ->
{symbols, classes}
info["type"] == "class" ->
class_entries = class_entries_for(library, python_module, info)
{symbols, class_entries ++ classes}
true ->
symbol_entry = symbol_entry_for(library, python_module, info)
{[symbol_entry | symbols], classes}
end
end)
end
defp symbol_entry_for(library, python_module, info) do
module = module_for_python(library, python_module)
python_name = info["python_name"] || info["name"]
{elixir_name, _python_name} = sanitize_function_name(python_name)
attribute? = info["type"] == "attribute"
params = info["parameters"] || []
{arity, arity_info} =
if attribute? do
{0, module_attr_arity_info()}
else
{required_arity(params), compute_arity_info(params, info)}
end
key = Manifest.symbol_key({module, String.to_atom(elixir_name), arity})
{
key,
%{
"module" => Module.split(module) |> Enum.join("."),
"function" => python_name,
"name" => elixir_name,
"python_name" => python_name,
"elixir_name" => elixir_name,
"python_module" => python_module,
"signature_available" => Map.get(info, "signature_available", true),
"parameters" => params,
"docstring" => info["docstring"] || "",
"return_annotation" => info["return_annotation"],
"return_type" => info["return_type"]
}
|> Map.merge(arity_info)
|> maybe_put_call_type(attribute?)
}
end
defp maybe_put_call_type(entry, true), do: Map.put(entry, "call_type", "module_attr")
defp maybe_put_call_type(entry, false), do: entry
defp module_attr_arity_info do
%{
"required_arity" => 0,
"minimum_arity" => 0,
"maximum_arity" => 0,
"has_var_positional" => false,
"has_var_keyword" => false,
"required_keyword_only" => [],
"optional_keyword_only" => []
}
end
defp class_entries_for(library, python_module, info) do
class_name = info["name"] || info["class"] || "Class"
class_python_module = info["python_module"] || python_module || library.python_name
class_module = class_module_for(library, class_python_module, class_name)
key = Manifest.class_key(class_module)
methods =
info["methods"]
|> List.wrap()
|> Enum.map(&class_method_entry/1)
|> Enum.reject(&is_nil/1)
[
{
key,
%{
"module" => Module.split(class_module) |> Enum.join("."),
"class" => class_name,
"python_module" => class_python_module,
"docstring" => info["docstring"] || "",
"methods" => methods,
"attributes" => info["attributes"] || []
}
}
]
end
defp class_method_entry(method) do
name = method["name"] || method[:name] || ""
case Generator.sanitize_method_name(name) do
{elixir_name, python_name} ->
params = method["parameters"] || method[:parameters] || []
arity_info = compute_arity_info(params, method)
arity_info =
if python_name == "__init__" do
arity_info
else
add_ref_arity_info(arity_info)
end
method
|> Map.put("name", python_name)
|> Map.put("python_name", python_name)
|> Map.put("elixir_name", elixir_name)
|> Map.merge(arity_info)
nil ->
nil
end
end
defp add_ref_arity_info(arity_info) do
min_arity = Map.get(arity_info, "minimum_arity", 0) + 1
required_arity = Map.get(arity_info, "required_arity", 0) + 1
max_arity = Map.get(arity_info, "maximum_arity")
max_arity =
case max_arity do
value when is_integer(value) -> value + 1
_ -> max_arity
end
arity_info
|> Map.put("minimum_arity", min_arity)
|> Map.put("required_arity", required_arity)
|> Map.put("maximum_arity", max_arity)
end
defp class_module_for(library, python_module, class_name) do
python_parts = String.split(python_module, ".")
library_parts = String.split(library.python_name, ".")
extra_parts = Enum.drop(python_parts, length(library_parts))
extra_parts = drop_class_suffix(extra_parts, class_name)
library.module_name
|> Module.split()
|> Kernel.++(Enum.map(extra_parts, &Macro.camelize/1))
|> Kernel.++([class_name])
|> Module.concat()
end
defp drop_class_suffix(parts, class_name) when is_list(parts) and is_binary(class_name) do
class_suffix = Macro.underscore(class_name)
case List.last(parts) do
^class_suffix -> Enum.drop(parts, -1)
_ -> parts
end
end
defp drop_class_suffix(parts, _class_name), do: parts
defp module_for_python(library, python_module) do
python_parts = String.split(python_module, ".")
library_parts = String.split(library.python_name, ".")
extra_parts = Enum.drop(python_parts, length(library_parts))
library.module_name
|> Module.split()
|> Kernel.++(Enum.map(extra_parts, &Macro.camelize/1))
|> Module.concat()
end
defp python_module_for_elixir(library, module) do
module_parts = Module.split(module)
library_parts = Module.split(library.module_name)
extra_parts = Enum.drop(module_parts, length(library_parts))
case Enum.map(extra_parts, &Macro.underscore/1) do
[] -> library.python_name
parts -> library.python_name <> "." <> Enum.join(parts, ".")
end
end
defp function_or_class_present?(manifest, library, name) do
module_prefix = Module.split(library.module_name) |> Enum.join(".")
function_present_in_manifest?(manifest, module_prefix, name) or
class_present_in_manifest?(manifest, name)
end
defp function_present_in_manifest?(manifest, module_prefix, name) do
manifest
|> Map.get("symbols", %{})
|> Map.values()
|> Enum.any?(&symbol_matches?(&1, module_prefix, name))
end
defp symbol_matches?(info, module_prefix, name) do
module = info["module"] || ""
python_name = info["python_name"] || info["function"] || info["name"]
elixir_name = info["elixir_name"] || info["name"] || python_name
module == module_prefix and (python_name == name or elixir_name == name)
end
defp class_present_in_manifest?(manifest, name) do
manifest
|> Map.get("classes", %{})
|> Map.values()
|> Enum.any?(&class_matches?(&1, name))
end
defp class_matches?(info, name) do
info["class"] == name or String.ends_with?(info["module"] || "", ".#{name}")
end
defp generate_from_manifest(config, manifest) do
Enum.each(config.libraries, fn library ->
functions = functions_for_library(manifest, library)
classes = classes_for_library(manifest, library)
Generator.generate_library(library, functions, classes, config)
end)
end
defp functions_for_library(manifest, library) do
Map.get(manifest, "symbols", %{})
|> Map.values()
|> Enum.filter(fn info ->
python_module = info["python_module"] || ""
String.starts_with?(python_module, library.python_name)
end)
end
defp classes_for_library(manifest, library) do
Map.get(manifest, "classes", %{})
|> Map.values()
|> Enum.filter(fn info ->
python_module = info["python_module"] || ""
String.starts_with?(python_module, library.python_name)
end)
end
defp strict_mode?(config) do
System.get_env("SNAKEBRIDGE_STRICT") == "1" || config.strict == true
end
defp run_strict(config) do
manifest = Manifest.load(config)
detected = scanner_module().scan_project(config)
missing = Manifest.missing(manifest, detected)
if missing != [] do
formatted = format_missing(missing)
raise SnakeBridge.CompileError, """
Strict mode: #{length(missing)} symbol(s) not in manifest.
Missing:
#{formatted}
To fix:
1. Run `mix snakebridge.setup` locally
2. Run `mix compile` to generate bindings
3. Commit the updated manifest and generated files
4. Re-run CI
Set SNAKEBRIDGE_STRICT=0 to disable strict mode.
"""
end
verify_generated_files_exist!(config)
verify_symbols_present!(config, manifest)
{:ok, []}
end
defp run_normal(config) do
start_time = System.monotonic_time()
libraries = Enum.map(config.libraries, & &1.name)
Telemetry.compile_start(libraries, false)
try do
detected = scanner_module().scan_project(config)
manifest = Manifest.load(config)
missing = Manifest.missing(manifest, detected)
targets = build_targets(missing, config, manifest)
updated_manifest =
if targets != [] do
update_manifest(manifest, targets)
else
manifest
end
Manifest.save(config, updated_manifest)
generate_from_manifest(config, updated_manifest)
generate_helper_wrappers(config)
SnakeBridge.Registry.save()
Lock.update(config)
symbol_count = count_symbols(updated_manifest)
file_count = length(config.libraries)
Telemetry.compile_stop(start_time, symbol_count, file_count, libraries, :normal)
{:ok, []}
rescue
e ->
Telemetry.compile_exception(start_time, e, __STACKTRACE__)
reraise e, __STACKTRACE__
end
end
@spec verify_generated_files_exist!(Config.t()) :: :ok
def verify_generated_files_exist!(config) do
Enum.each(config.libraries, fn library ->
path = Path.join(config.generated_dir, "#{library.python_name}.ex")
unless File.exists?(path) do
raise SnakeBridge.CompileError, """
Strict mode: Generated file missing: #{path}
Run `mix compile` locally and commit the generated files.
"""
end
end)
:ok
end
@spec verify_symbols_present!(Config.t(), map()) :: :ok
def verify_symbols_present!(config, manifest) do
Enum.each(config.libraries, fn library ->
path = Path.join(config.generated_dir, "#{library.python_name}.ex")
content = read_generated_file!(path)
defs = parse_definitions!(content, path)
missing_functions = missing_functions_for_library(manifest, library, defs)
classes_for_library = classes_for_library(manifest, library)
missing_classes = missing_classes_for_library(classes_for_library, defs)
missing_class_members = missing_class_members_for_library(classes_for_library, defs)
maybe_raise_missing!(path, missing_functions, missing_classes, missing_class_members)
end)
:ok
end
defp missing_functions_for_library(manifest, library, defs) do
manifest
|> functions_for_library(library)
|> Enum.reject(fn info ->
function_defined?(defs, info["module"], info["name"])
end)
end
defp missing_classes_for_library(classes_for_library, defs) do
Enum.reject(classes_for_library, fn info ->
module = info["module"]
Map.has_key?(defs, module)
end)
end
defp missing_class_members_for_library(classes_for_library, defs) do
classes_for_library
|> Enum.filter(fn info -> Map.has_key?(defs, info["module"]) end)
|> Enum.flat_map(&missing_members_for_class(&1, defs))
end
defp missing_members_for_class(info, defs) do
module = info["module"]
method_names = missing_method_names(info["methods"] || [], defs, module)
attr_names = missing_attr_names(info["attributes"] || [], defs, module)
Enum.map(method_names ++ attr_names, fn name ->
{module, name}
end)
end
defp missing_method_names(methods, defs, module) do
methods
|> Enum.map(&method_expected_name/1)
|> Enum.reject(fn name ->
name == "" or function_defined?(defs, module, name)
end)
end
defp missing_attr_names(attrs, defs, module) do
attrs
|> Enum.map(&to_string/1)
|> Enum.reject(fn name ->
name == "" or function_defined?(defs, module, name)
end)
end
defp maybe_raise_missing!(path, missing_functions, missing_classes, missing_class_members) do
if missing_functions != [] or missing_classes != [] or missing_class_members != [] do
raise SnakeBridge.CompileError, """
Strict mode: Generated file #{path} is missing expected bindings.
#{missing_functions_message(missing_functions)}\
#{missing_classes_message(missing_classes)}\
#{missing_class_members_message(missing_class_members)}
Run `mix compile` locally to regenerate and commit the updated files.
"""
end
end
defp required_arity(params) do
params
|> Enum.filter(fn param ->
param_kind(param) in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"]
end)
|> Enum.reject(¶m_default?/1)
|> length()
end
defp compute_arity_info(params, info) do
required_positional = required_arity(params)
optional_positional = params |> Enum.filter(&optional_positional?/1) |> length()
has_var_positional = Enum.any?(params, &varargs?/1)
has_var_keyword = Enum.any?(params, &kwargs?/1)
required_kw_only =
params
|> Enum.filter(&keyword_only_required?/1)
|> Enum.map(& &1["name"])
optional_kw_only =
params
|> Enum.filter(&keyword_only_optional?/1)
|> Enum.map(& &1["name"])
signature_available = Map.get(info, "signature_available", true)
variadic_fallback = params == [] and signature_available == false
max_arity =
cond do
variadic_fallback -> variadic_max_arity() + 1
has_var_positional -> :unbounded
optional_positional > 0 -> required_positional + 2
true -> required_positional + 1
end
%{
"required_arity" => required_positional,
"minimum_arity" => required_positional,
"maximum_arity" => max_arity,
"has_var_positional" => has_var_positional,
"has_var_keyword" => has_var_keyword,
"required_keyword_only" => required_kw_only,
"optional_keyword_only" => optional_kw_only
}
end
defp optional_positional?(param) do
param_kind(param) in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"] and param_default?(param)
end
defp varargs?(param), do: param_kind(param) == "VAR_POSITIONAL"
defp kwargs?(param), do: param_kind(param) == "VAR_KEYWORD"
defp keyword_only_required?(param) do
param_kind(param) == "KEYWORD_ONLY" and not param_default?(param)
end
defp keyword_only_optional?(param) do
param_kind(param) == "KEYWORD_ONLY" and param_default?(param)
end
defp param_kind(%{"kind" => kind}) when is_binary(kind), do: String.upcase(kind)
defp param_kind(%{kind: kind}) when is_binary(kind), do: String.upcase(kind)
defp param_kind(%{kind: kind}), do: kind
defp param_kind(_), do: nil
defp param_default?(%{"default" => _}), do: true
defp param_default?(%{default: _}), do: true
defp param_default?(_), do: false
defp variadic_max_arity do
Application.get_env(:snakebridge, :variadic_max_arity, 8)
end
defp sanitize_function_name(python_name) when is_binary(python_name) do
elixir_name =
python_name
|> Macro.underscore()
|> String.replace(~r/[^a-z0-9_?!]/, "_")
|> ensure_valid_identifier()
elixir_name =
if elixir_name in @reserved_words do
"py_#{elixir_name}"
else
elixir_name
end
{elixir_name, python_name}
end
defp ensure_valid_identifier(""), do: "_"
defp ensure_valid_identifier(name) do
if String.match?(name, ~r/^[a-z_][a-z0-9_?!]*$/) do
name
else
"_" <> name
end
end
defp python_name_for_elixir_name(elixir_name) when is_binary(elixir_name) do
case String.split(elixir_name, "py_", parts: 2) do
["", rest] when rest in @reserved_words -> rest
_ -> elixir_name
end
end
defp format_missing(missing) do
missing
|> Enum.sort()
|> Enum.map_join("\n", fn {mod, fun, arity} ->
module = Module.split(mod) |> Enum.join(".")
" - #{module}.#{fun}/#{arity}"
end)
end
defp scanner_module do
Application.get_env(:snakebridge, :scanner, Scanner)
end
defp read_generated_file!(path) do
case File.read(path) do
{:ok, content} ->
content
{:error, reason} ->
raise SnakeBridge.CompileError, """
Strict mode: Cannot read generated file #{path}: #{inspect(reason)}
"""
end
end
defp parse_definitions!(content, path) do
case Code.string_to_quoted(content, file: path) do
{:ok, ast} ->
collect_definitions(ast)
{:error, {line, error, token}} ->
raise SnakeBridge.CompileError, """
Strict mode: Cannot parse generated file #{path}: #{error} #{inspect(token)} on line #{line}
"""
end
end
defp collect_definitions(ast) do
initial = %{stack: [], defs: %{}}
{_, acc} =
Macro.traverse(ast, initial, &collect_pre/2, &collect_post/2)
acc.defs
end
defp collect_pre({:defmodule, _, [{:__aliases__, _, parts}, _]} = node, acc) do
segments = module_segments(parts, acc.stack)
module_name = Enum.join(segments, ".")
acc =
acc
|> Map.update!(:stack, &[segments | &1])
|> Map.update(:defs, %{}, &Map.put_new(&1, module_name, MapSet.new()))
{node, acc}
end
defp collect_pre({:def, _, [head | _]} = node, acc) do
{node, track_def(acc, head)}
end
defp collect_pre(node, acc), do: {node, acc}
defp track_def(acc, head) do
case {def_name(head), List.first(acc.stack)} do
{nil, _} ->
acc
{_, nil} ->
acc
{name, current_module} ->
module_name = Enum.join(current_module, ".")
Map.update(acc, :defs, %{}, fn defs ->
Map.update(defs, module_name, MapSet.new([name]), &MapSet.put(&1, name))
end)
end
end
defp collect_post({:defmodule, _, _} = node, acc) do
{_popped, rest} = List.pop_at(acc.stack, 0)
{node, %{acc | stack: rest}}
end
defp collect_post(node, acc), do: {node, acc}
defp module_segments(parts, []) do
Enum.map(parts, &Atom.to_string/1)
end
defp module_segments(parts, [parent | _]) do
parent ++ Enum.map(parts, &Atom.to_string/1)
end
defp def_name({:when, _, [inner | _]}), do: def_name(inner)
defp def_name({name, _, _}) when is_atom(name), do: Atom.to_string(name)
defp def_name(_), do: nil
defp function_defined?(defs, module, name) when is_binary(module) and is_binary(name) do
case Map.get(defs, module) do
nil -> false
set -> MapSet.member?(set, name)
end
end
defp method_expected_name(%{"elixir_name" => name}) when is_binary(name), do: name
defp method_expected_name(%{elixir_name: name}) when is_binary(name), do: name
defp method_expected_name(%{"name" => name}) when is_binary(name) do
case Generator.sanitize_method_name(name) do
{elixir_name, _} -> elixir_name
_ -> ""
end
end
defp method_expected_name(%{name: name}) when is_binary(name) do
case Generator.sanitize_method_name(name) do
{elixir_name, _} -> elixir_name
_ -> ""
end
end
defp method_expected_name(_), do: ""
defp missing_functions_message([]), do: ""
defp missing_functions_message(missing) do
formatted =
missing
|> Enum.map_join("\n", fn info ->
module = info["module"] || "Unknown"
name = info["name"] || "unknown"
" - #{module}.#{name}"
end)
"""
Missing functions:
#{formatted}
"""
end
defp missing_classes_message([]), do: ""
defp missing_classes_message(missing) do
formatted =
missing
|> Enum.map_join("\n", fn info ->
info["module"] || "Unknown"
end)
"""
Missing classes:
#{formatted}
"""
end
defp missing_class_members_message([]), do: ""
defp missing_class_members_message(missing) do
formatted =
missing
|> Enum.map_join("\n", fn {module, name} ->
" - #{module}.#{name}"
end)
"""
Missing class members:
#{formatted}
"""
end
defp count_symbols(manifest) do
symbols = Map.get(manifest, "symbols", %{}) |> map_size()
classes = Map.get(manifest, "classes", %{}) |> map_size()
symbols + classes
end
defp generate_helper_wrappers(config) do
if Helpers.enabled?(config) do
case Helpers.discover(config) do
{:ok, helpers} ->
HelperGenerator.generate_helpers(helpers, config)
{:error, %SnakeBridge.HelperRegistryError{} = error} ->
Mix.shell().error(Exception.message(error))
:ok
{:error, reason} ->
Mix.shell().error("Helper registry failed: #{inspect(reason)}")
:ok
end
else
:ok
end
end
end
defmodule Mix.Tasks.Snakebridge.Setup do
@shortdoc "Provision Python environment for SnakeBridge"
@moduledoc """
Provisions the Python environment for SnakeBridge introspection.
## Usage
mix snakebridge.setup
## Options
--upgrade Upgrade packages to latest matching versions
--verbose Show detailed output
--check Only check, don't install (exit 1 if missing)
"""
use Mix.Task
def run(args) do
{opts, _, _} =
OptionParser.parse(args,
switches: [upgrade: :boolean, verbose: :boolean, check: :boolean]
)
Mix.Task.run("loadpaths")
config = SnakeBridge.Config.load()
requirements = SnakeBridge.PythonEnv.derive_requirements(config.libraries)
if requirements == [] do
Mix.shell().info("No Python packages required (all stdlib)")
:ok
else
if opts[:check] do
run_check(requirements)
else
ensure_python_runtime!()
run_install(requirements, opts)
end
end
end
defp run_check(requirements) do
case python_packages_module().check_installed(requirements, python_packages_opts([])) do
{:ok, :all_installed} ->
Mix.shell().info("All packages installed")
{:ok, {:missing, missing}} ->
Mix.raise("Missing packages: #{inspect(missing)}")
end
end
defp run_install(requirements, opts) do
Mix.shell().info("Installing Python packages...")
install_opts = [
upgrade: opts[:upgrade] || false,
quiet: !opts[:verbose]
]
python_packages_module().ensure!({:list, requirements}, python_packages_opts(install_opts))
Mix.shell().info("Done. #{length(requirements)} package(s) ready.")
end
defp ensure_python_runtime! do
python_config = Application.get_env(:snakepit, :python, [])
if Keyword.get(python_config, :managed, false) do
python_runtime_module().install_managed(SnakeBridge.PythonRuntimeRunner, [])
end
:ok
end
defp python_packages_module do
Application.get_env(:snakebridge, :python_packages, Snakepit.PythonPackages)
end
defp python_packages_opts(opts) do
if python_packages_module() == Snakepit.PythonPackages do
Keyword.put_new(opts, :runner, SnakeBridge.PythonPackagesRunner)
else
opts
end
end
defp python_runtime_module do
Application.get_env(:snakebridge, :python_runtime, Snakepit.PythonRuntime)
end
end
defmodule Mix.Tasks.Snakebridge.Verify do
@moduledoc """
Verifies the lock file against the current hardware environment.
This task checks that the hardware environment where the lock file was created
is compatible with the current system. It detects:
- Platform mismatches (OS, architecture)
- CUDA version differences
- Missing GPU capabilities
- CPU feature mismatches
## Usage
mix snakebridge.verify # Verify with warnings
mix snakebridge.verify --strict # Fail on any mismatch
mix snakebridge.verify --verbose # Show detailed info
## Options
- `--strict` - Treat warnings as errors and fail
- `--verbose` - Print detailed hardware information
- `--file PATH` - Use a specific lock file (default: snakebridge.lock)
## Exit Codes
- 0 - Compatible environment
- 1 - Incompatible environment (or warnings in strict mode)
## Examples
# Standard verification
$ mix snakebridge.verify
✓ Lock file compatible with current environment
# Strict mode (CI)
$ mix snakebridge.verify --strict
✗ CUDA version mismatch: lock has 12.1, current has 11.8
# Verbose output
$ mix snakebridge.verify --verbose
Current hardware:
Platform: linux-x86_64
Accelerator: cuda
CUDA version: 12.1
GPU count: 2
Lock file:
Platform: linux-x86_64
Accelerator: cuda
CUDA version: 12.1
GPU count: 2
✓ Lock file compatible
"""
use Mix.Task
alias SnakeBridge.Lock.Verifier
@shortdoc "Verify lock file compatibility with current hardware"
@switches [
strict: :boolean,
verbose: :boolean,
file: :string
]
@impl Mix.Task
def run(args) do
{opts, _remaining, _invalid} = OptionParser.parse(args, switches: @switches)
lock_file = Keyword.get(opts, :file, "snakebridge.lock")
strict? = Keyword.get(opts, :strict, false)
verbose? = Keyword.get(opts, :verbose, false)
case load_lock(lock_file) do
{:ok, lock} ->
if verbose?, do: print_verbose_info(lock)
verify_and_report(lock, strict?)
{:error, :not_found} ->
Mix.shell().error("Lock file not found: #{lock_file}")
Mix.shell().error("Run `mix compile` to generate the lock file.")
raise Mix.Error, message: "Lock file not found"
{:error, reason} ->
Mix.shell().error("Failed to load lock file: #{inspect(reason)}")
raise Mix.Error, message: "Failed to load lock file"
end
end
defp load_lock(path) do
case File.read(path) do
{:ok, content} ->
{:ok, Jason.decode!(content)}
{:error, :enoent} ->
{:error, :not_found}
{:error, reason} ->
{:error, reason}
end
end
defp verify_and_report(lock, strict?) do
case Verifier.verify(lock) do
:ok ->
Mix.shell().info("✓ Lock file compatible with current environment")
:ok
{:warning, warnings} ->
Enum.each(warnings, fn warning ->
Mix.shell().info("⚠ Warning: #{warning}")
end)
if strict? do
Mix.shell().error("Strict mode: treating warnings as errors")
raise Mix.Error, message: "Lock file has compatibility warnings"
else
Mix.shell().info("✓ Lock file compatible (with warnings)")
:ok
end
{:error, errors} ->
Enum.each(errors, fn error ->
Mix.shell().error("✗ Error: #{error}")
end)
raise Mix.Error, message: "Lock file incompatible with current environment"
end
end
defp print_verbose_info(lock) do
current = hardware_module().identity()
current_caps = hardware_module().capabilities()
Mix.shell().info("Current hardware:")
Mix.shell().info(" Platform: #{current["platform"]}")
Mix.shell().info(" Accelerator: #{current["accelerator"]}")
if current_caps.cuda do
Mix.shell().info(" CUDA version: #{current_caps.cuda_version}")
end
if current_caps.mps do
Mix.shell().info(" MPS: available")
end
Mix.shell().info(" GPU count: #{current["gpu_count"]}")
Mix.shell().info(" CPU features: #{Enum.join(current["cpu_features"] || [], ", ")}")
Mix.shell().info("")
lock_env = Map.get(lock, "environment", %{})
lock_hardware = Map.get(lock_env, "hardware", %{})
lock_platform = Map.get(lock_env, "platform", %{})
Mix.shell().info("Lock file:")
Mix.shell().info(" Platform: #{lock_platform["os"]}-#{lock_platform["arch"]}")
Mix.shell().info(" Accelerator: #{lock_hardware["accelerator"]}")
if cuda_version = lock_hardware["cuda_version"] do
Mix.shell().info(" CUDA version: #{cuda_version}")
end
Mix.shell().info(" GPU count: #{lock_hardware["gpu_count"]}")
cpu_features = lock_hardware["cpu_features"] || []
Mix.shell().info(" CPU features: #{Enum.join(cpu_features, ", ")}")
Mix.shell().info("")
end
defp hardware_module do
Application.get_env(:snakebridge, :hardware_module, Snakepit.Hardware)
end
end
defmodule SnakeBridge.Docs.MarkdownConverter do
@moduledoc """
Converts parsed Python docstrings to Elixir ExDoc Markdown format.
This module transforms structured docstring data into Markdown that
is compatible with ExDoc and follows Elixir documentation conventions.
"""
alias SnakeBridge.Docs.MathRenderer
@type_map %{
"int" => "integer()",
"float" => "float()",
"str" => "String.t()",
"string" => "String.t()",
"bool" => "boolean()",
"boolean" => "boolean()",
"None" => "nil",
"NoneType" => "nil",
"list" => "list()",
"dict" => "map()",
"tuple" => "tuple()",
"set" => "MapSet.t()",
"bytes" => "binary()",
"bytearray" => "binary()",
"Any" => "term()",
"object" => "term()"
}
@exception_map %{
"ValueError" => "ArgumentError",
"TypeError" => "ArgumentError",
"KeyError" => "KeyError",
"IndexError" => "Enum.OutOfBoundsError",
"RuntimeError" => "RuntimeError",
"NotImplementedError" => "RuntimeError",
"IOError" => "File.Error",
"OSError" => "File.Error",
"FileNotFoundError" => "File.Error",
"AttributeError" => "KeyError",
"NameError" => "UndefinedFunctionError"
}
# Section builders - each returns nil if the section should be skipped
@section_builders [
:short_description,
:long_description,
:params,
:returns,
:raises,
:examples,
:notes
]
@doc """
Converts a parsed docstring structure to ExDoc Markdown format.
## Parameters
- `parsed` - A map with keys: `:short_description`, `:long_description`,
`:params`, `:returns`, `:raises`, `:examples`
## Returns
A Markdown string suitable for use in `@doc` or `@moduledoc`.
"""
@spec convert(map()) :: String.t()
def convert(parsed) when is_map(parsed) do
@section_builders
|> Enum.map(&build_section(&1, parsed))
|> Enum.reject(&is_nil/1)
|> Enum.join("\n\n")
|> MathRenderer.render()
|> String.trim()
end
defp build_section(:short_description, parsed), do: parsed[:short_description]
defp build_section(:long_description, parsed), do: parsed[:long_description]
defp build_section(:params, %{params: params}) when is_list(params) and params != [] do
format_parameters(params)
end
defp build_section(:params, _parsed), do: nil
defp build_section(:returns, %{returns: returns}) when not is_nil(returns) do
format_returns(returns)
end
defp build_section(:returns, _parsed), do: nil
defp build_section(:raises, %{raises: raises}) when is_list(raises) and raises != [] do
format_raises(raises)
end
defp build_section(:raises, _parsed), do: nil
defp build_section(:examples, %{examples: examples})
when is_list(examples) and examples != [] do
format_examples(examples)
end
defp build_section(:examples, _parsed), do: nil
defp build_section(:notes, %{notes: notes}) when not is_nil(notes) do
"## Notes\n\n#{notes}"
end
defp build_section(:notes, _parsed), do: nil
# Generic type patterns with their prefixes and converters
# Format: {prefix, prefix_length, converter_function}
@generic_type_patterns [
{"Optional[", 9, :convert_optional},
{"Union[", 6, :convert_union},
{"list[", 5, :convert_list},
{"List[", 5, :convert_list},
{"dict[", 5, :convert_dict},
{"Dict[", 5, :convert_dict},
{"tuple[", 6, :convert_tuple},
{"Tuple[", 6, :convert_tuple},
{"set[", 4, :convert_set},
{"Set[", 4, :convert_set}
]
@doc """
Converts a Python type annotation to an Elixir typespec format.
## Examples
iex> MarkdownConverter.convert_type("int")
"integer()"
iex> MarkdownConverter.convert_type("list[str]")
"list(String.t())"
"""
@spec convert_type(String.t() | nil) :: String.t()
def convert_type(nil), do: "term()"
def convert_type(""), do: "term()"
def convert_type(python_type) do
python_type = String.trim(python_type)
case Map.fetch(@type_map, python_type) do
{:ok, elixir_type} -> elixir_type
:error -> convert_generic_type(python_type)
end
end
defp convert_generic_type(python_type) do
@generic_type_patterns
|> Enum.find_value(fn {prefix, prefix_len, converter} ->
if String.starts_with?(python_type, prefix) do
inner = extract_inner_type(python_type, prefix_len)
apply_type_converter(converter, inner)
end
end)
|> Kernel.||(python_type)
end
defp extract_inner_type(type_string, prefix_length) do
String.slice(type_string, prefix_length..-2//1)
end
defp apply_type_converter(:convert_optional, inner) do
"#{convert_type(inner)} | nil"
end
defp apply_type_converter(:convert_union, inner) do
inner
|> String.split(",")
|> Enum.map_join(" | ", &(&1 |> String.trim() |> convert_type()))
end
defp apply_type_converter(:convert_list, inner) do
"list(#{convert_type(inner)})"
end
defp apply_type_converter(:convert_dict, _inner) do
"map()"
end
defp apply_type_converter(:convert_tuple, _inner) do
"tuple()"
end
defp apply_type_converter(:convert_set, inner) do
"MapSet.t(#{convert_type(inner)})"
end
@doc """
Converts a Python exception type to an Elixir exception module.
## Examples
iex> MarkdownConverter.convert_exception("ValueError")
"ArgumentError"
"""
@spec convert_exception(String.t() | nil) :: String.t()
def convert_exception(nil), do: "RuntimeError"
def convert_exception(""), do: "RuntimeError"
def convert_exception(python_exception) do
Map.get(@exception_map, python_exception, python_exception)
end
@doc """
Converts a Python doctest example to Elixir iex format.
## Examples
iex> MarkdownConverter.convert_example(">>> func(1, 2)\\n3")
" iex> func(1, 2)\\n 3"
"""
@spec convert_example(String.t()) :: String.t()
def convert_example(example) do
Enum.map_join(String.split(example, "\n"), "\n", fn line ->
cond do
String.starts_with?(String.trim(line), ">>>") ->
code = line |> String.trim() |> String.slice(3..-1//1) |> String.trim()
" iex> #{code}"
String.starts_with?(String.trim(line), "...") ->
code = line |> String.trim() |> String.slice(3..-1//1) |> String.trim()
" ...> #{code}"
String.trim(line) == "" ->
""
true ->
" #{String.trim(line)}"
end
end)
end
defp format_parameters(params) do
param_lines =
Enum.map(params, fn param ->
name = param[:name] || param.name
type_name = param[:type_name]
description = param[:description]
type_str =
if type_name do
" (type: `#{convert_type(type_name)}`)"
else
""
end
default_str =
if param[:default] do
" Defaults to `#{param.default}`."
else
""
end
desc_str = if description, do: " - #{description}", else: ""
"- `#{name}`#{desc_str}#{type_str}#{default_str}"
end)
"## Parameters\n\n#{Enum.join(param_lines, "\n")}"
end
defp format_returns(returns) do
type_str =
if returns[:type_name] do
"Returns `#{convert_type(returns.type_name)}`."
else
""
end
desc_str =
if returns[:description] do
" #{returns.description}"
else
""
end
"## Returns\n\n#{type_str}#{desc_str}"
end
defp format_raises(raises) do
raise_lines =
Enum.map(raises, fn r ->
type = convert_exception(r[:type_name] || r.type_name)
desc = r[:description] || ""
"- `#{type}` - #{desc}"
end)
"## Raises\n\n#{Enum.join(raise_lines, "\n")}"
end
defp format_examples(examples) do
formatted = Enum.map_join(examples, "\n\n", &convert_example/1)
"## Examples\n\n#{formatted}"
end
end
defmodule SnakeBridge.Docs.MathRenderer do
@moduledoc """
Renders LaTeX math expressions for documentation.
Converts reStructuredText math directives to Markdown-compatible
math notation (KaTeX/MathJax style).
## Supported Formats
- Inline math: ``:math:`E = mc^2` `` → `$E = mc^2$`
- Display math: `.. math::` blocks → `$$...$$`
"""
@doc """
Renders math expressions in a docstring, converting RST math to Markdown.
## Examples
iex> MathRenderer.render("The formula is :math:`E = mc^2`.")
"The formula is $E = mc^2$."
"""
@spec render(String.t() | nil) :: String.t() | nil
def render(nil), do: nil
def render(""), do: ""
def render(text) when is_binary(text) do
text
|> render_inline_math()
|> render_display_math()
end
@doc """
Extracts all math expressions from text.
Returns a list of math expression strings (without delimiters).
"""
@spec extract_math(String.t() | nil) :: [String.t()]
def extract_math(nil), do: []
def extract_math(""), do: []
def extract_math(text) do
inline = extract_inline_math(text)
display = extract_display_math(text)
inline ++ display
end
@doc """
Converts math expressions to KaTeX-compatible format.
KaTeX uses `$...$` for inline and `$$...$$` for display math.
"""
@spec to_katex(String.t() | nil) :: String.t() | nil
def to_katex(nil), do: nil
def to_katex(""), do: ""
def to_katex(text) do
# KaTeX format is the same as our render output
render(text)
end
# Render inline math: :math:`expr` → $expr$
defp render_inline_math(text) do
Regex.replace(
~r/:math:`([^`]+)`/,
text,
"$\\1$"
)
end
# Render display math blocks
defp render_display_math(text) do
# Pattern for RST math blocks:
# .. math::
#
# expression
Regex.replace(
~r/\.\.\s*math::\s*\n\s*\n((?:\s+.+\n?)+)/,
text,
fn _, content ->
expr =
content
|> String.trim()
|> String.replace(~r/^\s+/m, "")
"\n$$\n#{expr}\n$$\n"
end
)
end
# Extract inline math expressions
defp extract_inline_math(text) do
Regex.scan(~r/:math:`([^`]+)`/, text)
|> Enum.map(fn [_, expr] -> expr end)
end
# Extract display math expressions
defp extract_display_math(text) do
Regex.scan(
~r/\.\.\s*math::\s*\n\s*\n((?:\s+.+\n?)+)/,
text
)
|> Enum.map(fn [_, content] ->
content
|> String.trim()
|> String.replace(~r/^\s+/m, "")
end)
end
end
defmodule SnakeBridge.Docs.RstParser do
@moduledoc """
Parses Python docstrings in various formats (Google, NumPy, Sphinx, Epytext).
This module detects the docstring format and extracts structured information
including parameters, return values, exceptions, and examples.
## Supported Formats
- **Google style**: Uses `Args:`, `Returns:`, `Raises:` sections
- **NumPy style**: Uses underlined section headers (`Parameters\n----------`)
- **Sphinx/reST style**: Uses `:param:`, `:type:`, `:returns:` directives
- **Epytext style**: Uses `@param`, `@type`, `@return` tags
"""
@type parsed_doc :: %{
short_description: String.t() | nil,
long_description: String.t() | nil,
params: [param()],
returns: returns() | nil,
raises: [raises()],
examples: [String.t()],
notes: String.t() | nil,
style: atom()
}
@type param :: %{
name: String.t(),
type_name: String.t() | nil,
description: String.t() | nil,
optional: boolean(),
default: String.t() | nil
}
@type returns :: %{
type_name: String.t() | nil,
description: String.t() | nil
}
@type raises :: %{
type_name: String.t(),
description: String.t() | nil
}
@doc """
Parses a Python docstring and returns structured data.
"""
@spec parse(String.t() | nil) :: parsed_doc()
def parse(nil), do: empty_result()
def parse(""), do: empty_result()
def parse(docstring) when is_binary(docstring) do
style = detect_style(docstring)
lines = String.split(docstring, "\n")
{short_desc, rest} = extract_short_description(lines)
{long_desc, sections} = extract_long_description(rest)
%{
short_description: short_desc,
long_description: long_desc,
params: extract_params(sections, style),
returns: extract_returns(sections, style),
raises: extract_raises(sections, style),
examples: extract_examples(sections, style),
notes: extract_notes(sections, style),
style: style
}
end
@doc """
Detects the docstring style based on content patterns.
"""
@spec detect_style(String.t() | nil) :: atom()
def detect_style(nil), do: :unknown
def detect_style(""), do: :unknown
def detect_style(docstring) do
cond do
epytext_style?(docstring) -> :epytext
sphinx_style?(docstring) -> :sphinx
numpy_style?(docstring) -> :numpy
google_style?(docstring) -> :google
true -> :unknown
end
end
defp epytext_style?(docstring) do
docstring =~ ~r/@param\s/ or docstring =~ ~r/@type\s/
end
defp sphinx_style?(docstring) do
docstring =~ ~r/:param\s+\w+:/ or docstring =~ ~r/:returns:/
end
defp numpy_style?(docstring) do
docstring =~ ~r/Parameters\n-+/ or docstring =~ ~r/Returns\n-+/
end
defp google_style?(docstring) do
docstring =~ ~r/\n\s*Args:\s*\n/ or
docstring =~ ~r/\n\s*Arguments:\s*\n/ or
docstring =~ ~r/\n\s*Returns:\s*\n/ or
docstring =~ ~r/\n\s*Raises:\s*\n/ or
docstring =~ ~r/\n\s*Example:\s*\n/ or
docstring =~ ~r/\n\s*Examples:\s*\n/
end
defp empty_result do
%{
short_description: nil,
long_description: nil,
params: [],
returns: nil,
raises: [],
examples: [],
notes: nil,
style: :unknown
}
end
defp extract_short_description([]), do: {nil, []}
defp extract_short_description(lines) do
# Skip leading empty lines
lines = Enum.drop_while(lines, &(String.trim(&1) == ""))
case lines do
[] ->
{nil, []}
[first | rest] ->
short = String.trim(first)
if short == "" do
{nil, rest}
else
{short, rest}
end
end
end
defp extract_long_description(lines) do
# Skip empty lines after short description
lines = Enum.drop_while(lines, &(String.trim(&1) == ""))
# Find where sections start
section_start = find_section_start(lines)
case section_start do
nil ->
# No sections, all is long description
long_desc =
lines
|> Enum.join("\n")
|> String.trim()
{(long_desc == "" && nil) || long_desc, []}
index ->
{desc_lines, section_lines} = Enum.split(lines, index)
long_desc =
desc_lines
|> Enum.join("\n")
|> String.trim()
{(long_desc == "" && nil) || long_desc, section_lines}
end
end
defp find_section_start(lines) do
Enum.find_index(lines, fn line ->
trimmed = String.trim(line)
# Google style sections
# NumPy style sections (check next line for dashes)
# Sphinx style
# Epytext style
trimmed in [
"Args:",
"Arguments:",
"Returns:",
"Yields:",
"Raises:",
"Example:",
"Examples:",
"Note:",
"Notes:",
"Warning:",
"Warnings:"
] or
trimmed in [
"Parameters",
"Returns",
"Yields",
"Raises",
"Examples",
"Notes",
"Warnings",
"See Also",
"References"
] or
String.starts_with?(trimmed, ":param ") or
String.starts_with?(trimmed, ":returns:") or
String.starts_with?(trimmed, "@param ") or
String.starts_with?(trimmed, "@return")
end)
end
defp extract_params(sections, :google), do: extract_google_params(sections)
defp extract_params(sections, :numpy), do: extract_numpy_params(sections)
defp extract_params(sections, :sphinx), do: extract_sphinx_params(sections)
defp extract_params(sections, :epytext), do: extract_epytext_params(sections)
defp extract_params(_sections, _style), do: []
defp extract_google_params(lines) do
lines
|> extract_section(["Args:", "Arguments:"])
|> parse_google_items()
|> Enum.map(&parse_google_param/1)
end
defp extract_numpy_params(lines) do
lines
|> extract_numpy_section("Parameters")
|> parse_numpy_items()
|> Enum.map(&parse_numpy_param/1)
end
defp extract_sphinx_params(lines) do
lines
|> Enum.filter(&String.starts_with?(String.trim(&1), ":param "))
|> Enum.map(&parse_sphinx_param/1)
end
defp extract_epytext_params(lines) do
lines
|> Enum.filter(&String.starts_with?(String.trim(&1), "@param "))
|> Enum.map(&parse_epytext_param/1)
end
defp parse_google_param({name_type, description}) do
{name, type_name, optional, default} = parse_param_name_type(name_type)
%{
name: name,
type_name: type_name,
description: description,
optional: optional,
default: default
}
end
defp parse_numpy_param({name_type, description}) do
[name_part | type_parts] = String.split(name_type, " : ", parts: 2)
name = String.trim(name_part)
type_info = if type_parts != [], do: hd(type_parts), else: nil
{type_name, optional} = parse_numpy_type(type_info)
%{
name: name,
type_name: type_name,
description: description,
optional: optional,
default: nil
}
end
defp parse_sphinx_param(line) do
case Regex.run(~r/:param\s+(\w+):\s*(.*)/, String.trim(line)) do
[_, name, desc] ->
%{name: name, type_name: nil, description: desc, optional: false, default: nil}
_ ->
%{name: "", type_name: nil, description: "", optional: false, default: nil}
end
end
defp parse_epytext_param(line) do
case Regex.run(~r/@param\s+(\w+):\s*(.*)/, String.trim(line)) do
[_, name, desc] ->
%{name: name, type_name: nil, description: desc, optional: false, default: nil}
_ ->
%{name: "", type_name: nil, description: "", optional: false, default: nil}
end
end
defp parse_param_name_type(name_type) do
# Pattern: "name (type, optional): description" or "name (type): description"
case Regex.run(~r/^(\w+)\s*\(([^)]+)\)/, name_type) do
[_, name, type_info] ->
optional = String.contains?(type_info, "optional")
type_name = type_info |> String.replace(~r/,?\s*optional/, "") |> String.trim()
default =
case Regex.run(~r/Defaults? to [`']?([^`']+)[`']?/, name_type) do
[_, val] -> val
_ -> nil
end
{name, type_name, optional, default}
_ ->
# Just name, no type
name = name_type |> String.split() |> List.first() || ""
{name, nil, false, nil}
end
end
defp parse_numpy_type(nil), do: {nil, false}
defp parse_numpy_type(type_info) do
optional = String.contains?(type_info, "optional")
type_name = type_info |> String.replace(~r/,?\s*optional/, "") |> String.trim()
{type_name, optional}
end
defp extract_returns(sections, :google), do: extract_google_returns(sections)
defp extract_returns(sections, :numpy), do: extract_numpy_returns(sections)
defp extract_returns(sections, :sphinx), do: extract_sphinx_returns(sections)
defp extract_returns(sections, :epytext), do: extract_epytext_returns(sections)
defp extract_returns(_sections, _style), do: nil
defp extract_google_returns(lines) do
case extract_section(lines, ["Returns:", "Yields:"]) |> parse_google_items() do
[] ->
nil
[{type_desc, description} | _] ->
%{type_name: String.trim(type_desc), description: description}
end
end
defp extract_numpy_returns(lines) do
case extract_numpy_section(lines, "Returns") |> parse_numpy_items() do
[] ->
nil
[{type_name, description} | _] ->
%{type_name: type_name, description: description}
end
end
defp extract_sphinx_returns(lines) do
case Enum.find(lines, &String.contains?(&1, ":returns:")) do
nil ->
nil
line ->
case Regex.run(~r/:returns:\s*(.*)/, line) do
[_, desc] -> %{type_name: nil, description: desc}
_ -> nil
end
end
end
defp extract_epytext_returns(lines) do
case Enum.find(lines, &String.contains?(&1, "@return")) do
nil ->
nil
line ->
case Regex.run(~r/@return\s*:?\s*(.*)/, line) do
[_, desc] -> %{type_name: nil, description: desc}
_ -> nil
end
end
end
defp extract_raises(sections, :google), do: extract_google_raises(sections)
defp extract_raises(sections, :numpy), do: extract_numpy_raises(sections)
defp extract_raises(sections, :sphinx), do: extract_sphinx_raises(sections)
defp extract_raises(sections, :epytext), do: extract_epytext_raises(sections)
defp extract_raises(_sections, _style), do: []
defp extract_google_raises(lines) do
lines
|> extract_section(["Raises:"])
|> parse_google_items()
|> Enum.map(fn {type_name, description} ->
%{type_name: String.trim(type_name), description: description}
end)
end
defp extract_numpy_raises(lines) do
lines
|> extract_numpy_section("Raises")
|> parse_numpy_items()
|> Enum.map(fn {type_name, description} ->
%{type_name: type_name, description: description}
end)
end
defp extract_sphinx_raises(lines) do
lines
|> Enum.filter(&String.contains?(&1, ":raises"))
|> Enum.map(fn line ->
case Regex.run(~r/:raises?\s+(\w+):\s*(.*)/, line) do
[_, type, desc] -> %{type_name: type, description: desc}
_ -> %{type_name: "Error", description: ""}
end
end)
end
defp extract_epytext_raises(lines) do
lines
|> Enum.filter(&String.contains?(&1, "@raise"))
|> Enum.map(fn line ->
case Regex.run(~r/@raise\s+(\w+):\s*(.*)/, line) do
[_, type, desc] -> %{type_name: type, description: desc}
_ -> %{type_name: "Error", description: ""}
end
end)
end
defp extract_examples(sections, style) do
section_content =
case style do
:google -> extract_section(sections, ["Example:", "Examples:"])
:numpy -> extract_numpy_section(sections, "Examples")
_ -> []
end
section_content
|> Enum.join("\n")
|> String.trim()
|> case do
"" -> []
content -> [content]
end
end
defp extract_notes(sections, style) do
section_content =
case style do
:google -> extract_section(sections, ["Note:", "Notes:"])
:numpy -> extract_numpy_section(sections, "Notes")
_ -> []
end
section_content
|> Enum.join("\n")
|> String.trim()
|> case do
"" -> nil
content -> content
end
end
defp extract_section(lines, headers) do
start_idx =
Enum.find_index(lines, fn line ->
String.trim(line) in headers
end)
case start_idx do
nil -> []
idx -> extract_section_content(lines, idx)
end
end
defp extract_section_content(lines, idx) do
lines
|> Enum.drop(idx + 1)
|> Enum.take_while(&line_in_section?/1)
end
defp line_in_section?(line) do
trimmed = String.trim(line)
cond do
trimmed == "" -> true
section_header?(trimmed) -> false
indented_line?(line) -> true
true -> false
end
end
defp indented_line?(line) do
String.starts_with?(line, " ") or String.starts_with?(line, "\t")
end
defp extract_numpy_section(lines, header) do
start_idx =
Enum.find_index(lines, fn line ->
String.trim(line) == header
end)
case start_idx do
nil -> []
idx -> extract_numpy_section_content(lines, idx)
end
end
defp extract_numpy_section_content(lines, idx) do
lines
|> Enum.drop(idx + 2)
|> Enum.with_index(idx + 2)
|> Enum.take_while(fn {line, line_idx} ->
line_in_numpy_section?(lines, line, line_idx)
end)
|> Enum.map(fn {line, _idx} -> line end)
end
defp line_in_numpy_section?(lines, line, line_idx) do
trimmed = String.trim(line)
cond do
trimmed == "" -> true
numpy_section_header?(lines, line_idx) -> false
true -> true
end
end
defp section_header?(line) do
line in [
"Args:",
"Arguments:",
"Returns:",
"Yields:",
"Raises:",
"Example:",
"Examples:",
"Note:",
"Notes:",
"Warning:",
"Warnings:"
]
end
defp numpy_section_header?(lines, idx) when is_integer(idx) and idx >= 0 do
case Enum.at(lines, idx + 1) do
nil -> false
next_line -> String.trim(next_line) =~ ~r/^-+$/
end
end
defp numpy_section_header?(_lines, _idx), do: false
defp parse_google_items(lines) do
lines
|> Enum.chunk_while(nil, &google_chunk_reducer/2, &chunk_finalizer/1)
|> Enum.reject(&is_nil/1)
|> Enum.map(&split_google_item/1)
end
defp google_chunk_reducer(line, acc) do
trimmed = String.trim(line)
cond do
trimmed == "" ->
emit_chunk_on_empty(acc)
google_item_header?(trimmed) ->
emit_and_start_new_chunk(acc, trimmed)
acc != nil ->
{:cont, acc <> " " <> trimmed}
true ->
{:cont, nil}
end
end
defp google_item_header?(trimmed) do
String.match?(trimmed, ~r/^\w+(\s*\([^)]*\))?:/)
end
defp split_google_item(item) do
case String.split(item, ":", parts: 2) do
[name_type, desc] -> {String.trim(name_type), String.trim(desc)}
[name_type] -> {String.trim(name_type), ""}
end
end
defp parse_numpy_items(lines) do
lines
|> Enum.chunk_while(nil, &numpy_chunk_reducer/2, &chunk_finalizer/1)
|> Enum.reject(&is_nil/1)
|> Enum.map(&split_numpy_item/1)
end
defp numpy_chunk_reducer(line, acc) do
trimmed = String.trim(line)
cond do
trimmed == "" ->
emit_chunk_on_empty(acc)
numpy_item_header?(line) ->
emit_and_start_new_chunk(acc, line)
acc != nil ->
{:cont, acc <> " " <> trimmed}
true ->
{:cont, nil}
end
end
defp numpy_item_header?(line) do
String.contains?(line, " : ") or not String.starts_with?(line, " ")
end
defp split_numpy_item(item) do
case String.split(item, "\n", parts: 2) do
[first, rest] ->
{String.trim(first), String.trim(rest)}
[first] ->
split_numpy_single_line(first)
end
end
defp split_numpy_single_line(line) do
case String.split(line, " : ", parts: 2) do
[name, desc] -> {String.trim(name), String.trim(desc)}
[name] -> {String.trim(name), ""}
end
end
# Shared chunk helpers for both Google and NumPy parsers
defp emit_chunk_on_empty(nil), do: {:cont, nil}
defp emit_chunk_on_empty(acc), do: {:cont, acc, nil}
defp emit_and_start_new_chunk(nil, new_value), do: {:cont, new_value}
defp emit_and_start_new_chunk(acc, new_value), do: {:cont, acc, new_value}
defp chunk_finalizer(nil), do: {:cont, []}
defp chunk_finalizer(acc), do: {:cont, acc, []}
end
defmodule SnakeBridge.Error.DtypeMismatchError do
@moduledoc """
Error for tensor dtype incompatibilities.
Provides information about expected vs actual dtypes and
suggestions for converting between types.
## Examples
iex> error = %SnakeBridge.Error.DtypeMismatchError{
...> expected: :float32,
...> got: :float64,
...> operation: :matmul,
...> message: "Expected float32 but got float64"
...> }
iex> Exception.message(error)
"Dtype mismatch in matmul..."
"""
@type dtype :: :float16 | :float32 | :float64 | :int32 | :int64 | :bool | atom()
@type t :: %__MODULE__{
expected: dtype(),
got: dtype(),
operation: atom() | nil,
message: String.t(),
suggestion: String.t(),
python_traceback: String.t() | nil
}
defexception [
:expected,
:got,
:operation,
:python_traceback,
message: "Dtype mismatch",
suggestion: "Convert tensor to the expected dtype"
]
@impl Exception
def message(%__MODULE__{} = error) do
op_str = if error.operation, do: " in #{error.operation}", else: ""
"""
Dtype mismatch#{op_str}
Expected: #{format_dtype(error.expected)}
Got: #{format_dtype(error.got)}
#{error.message}
Suggestion: #{error.suggestion}
"""
|> String.trim()
end
@doc """
Creates a DtypeMismatchError error with conversion suggestion.
"""
@spec new(dtype(), dtype(), keyword()) :: t()
def new(expected, got, opts \\ []) do
suggestion =
Keyword.get(opts, :suggestion) ||
generate_suggestion(expected, got)
%__MODULE__{
expected: expected,
got: got,
operation: Keyword.get(opts, :operation),
message: Keyword.get(opts, :message, "Types do not match"),
suggestion: suggestion,
python_traceback: Keyword.get(opts, :python_traceback)
}
end
@doc """
Generates a suggestion for converting between dtypes.
"""
@spec generate_suggestion(dtype(), dtype()) :: String.t()
def generate_suggestion(expected, got) do
_from = format_dtype(got)
to = format_dtype(expected)
cond do
precision_loss?(got, expected) ->
"Convert with tensor.to(torch.#{to}) - note: this may lose precision"
requires_explicit?(got, expected) ->
"Convert with tensor.to(torch.#{to})"
true ->
"Use tensor.to(torch.#{to}) or tensor.type(torch.#{expected_torch_type(expected)})"
end
end
defp format_dtype(dtype) when is_atom(dtype) do
dtype
|> Atom.to_string()
|> String.replace("_", "")
end
defp format_dtype(dtype), do: inspect(dtype)
defp expected_torch_type(:float16), do: "HalfTensor"
defp expected_torch_type(:float32), do: "FloatTensor"
defp expected_torch_type(:float64), do: "DoubleTensor"
defp expected_torch_type(:int32), do: "IntTensor"
defp expected_torch_type(:int64), do: "LongTensor"
defp expected_torch_type(:bool), do: "BoolTensor"
defp expected_torch_type(other), do: Atom.to_string(other)
# Detect if conversion loses precision
defp precision_loss?(from, to) do
precision_rank(from) > precision_rank(to)
end
defp precision_rank(:float64), do: 3
defp precision_rank(:float32), do: 2
defp precision_rank(:float16), do: 1
defp precision_rank(:int64), do: 2
defp precision_rank(:int32), do: 1
defp precision_rank(_), do: 0
# Detect if explicit conversion is required (e.g., float to int)
defp requires_explicit?(from, to) do
float_type?(from) != float_type?(to)
end
defp float_type?(dtype) when dtype in [:float16, :float32, :float64], do: true
defp float_type?(_), do: false
end
defmodule SnakeBridge.Error.OutOfMemoryError do
@moduledoc """
GPU out-of-memory error with recovery suggestions.
Provides detailed information about memory failures including
device info, memory stats, and actionable suggestions.
## Examples
iex> error = %SnakeBridge.Error.OutOfMemoryError{
...> device: {:cuda, 0},
...> requested_mb: 8192,
...> available_mb: 2048,
...> message: "CUDA out of memory"
...> }
iex> Exception.message(error)
"GPU Out of Memory on CUDA:0..."
"""
@type device :: :cpu | {:cuda, non_neg_integer()} | :mps | atom()
@type t :: %__MODULE__{
device: device(),
requested_mb: non_neg_integer() | nil,
available_mb: non_neg_integer() | nil,
total_mb: non_neg_integer() | nil,
message: String.t(),
suggestions: [String.t()],
python_traceback: String.t() | nil
}
defexception [
:device,
:requested_mb,
:available_mb,
:total_mb,
:python_traceback,
message: "Out of memory",
suggestions: []
]
@impl Exception
def message(%__MODULE__{} = error) do
parts = ["GPU Out of Memory on #{format_device(error.device)}"]
parts =
if error.requested_mb || error.available_mb || error.total_mb do
mem_info = [
"Memory Info:",
" Requested: #{error.requested_mb || "unknown"} MB",
" Available: #{error.available_mb || "unknown"} MB",
" Total: #{error.total_mb || "unknown"} MB"
]
parts ++ [""] ++ mem_info
else
parts
end
suggestions =
(error.suggestions ++ default_suggestions(error.device))
|> Enum.uniq()
|> Enum.with_index(1)
|> Enum.map(fn {s, i} -> " #{i}. #{s}" end)
parts = parts ++ ["", "Suggestions:"] ++ suggestions
parts
|> Enum.join("\n")
|> String.trim()
end
@doc """
Creates an OutOfMemoryError error with default suggestions.
"""
@spec new(device(), keyword()) :: t()
def new(device, opts \\ []) do
%__MODULE__{
device: device,
requested_mb: Keyword.get(opts, :requested_mb),
available_mb: Keyword.get(opts, :available_mb),
total_mb: Keyword.get(opts, :total_mb),
message: Keyword.get(opts, :message, "Out of memory on #{format_device(device)}"),
suggestions: Keyword.get(opts, :suggestions, []),
python_traceback: Keyword.get(opts, :python_traceback)
}
end
defp format_device(:cpu), do: "CPU"
defp format_device(:mps), do: "Apple MPS"
defp format_device({:cuda, id}), do: "CUDA:#{id}"
defp format_device(other), do: inspect(other)
defp default_suggestions(device) do
base = [
"Reduce batch size",
"Use gradient checkpointing",
"Enable mixed precision training",
"Clear cached memory"
]
case device do
{:cuda, _} -> base ++ ["Move some operations to CPU"]
:mps -> base ++ ["Move some operations to CPU"]
_ -> base
end
end
end
defmodule SnakeBridge.Error.ShapeMismatchError do
@moduledoc """
Error for tensor shape incompatibilities.
This error provides detailed information about shape mismatches including
the operation that failed, the shapes involved, and actionable suggestions.
## Examples
iex> error = %SnakeBridge.Error.ShapeMismatchError{
...> operation: :matmul,
...> shape_a: [3, 4],
...> shape_b: [2, 5],
...> message: "Cannot multiply matrices with incompatible shapes",
...> suggestion: "A has 4 columns but B has 2 rows. Transpose B."
...> }
iex> Exception.message(error)
"Shape mismatch in matmul..."
"""
@type t :: %__MODULE__{
operation: atom(),
shape_a: [non_neg_integer()] | nil,
shape_b: [non_neg_integer()] | nil,
expected: String.t() | nil,
got: String.t() | nil,
message: String.t(),
suggestion: String.t(),
python_traceback: String.t() | nil
}
defexception [
:operation,
:shape_a,
:shape_b,
:expected,
:got,
:python_traceback,
message: "Shape mismatch",
suggestion: "Check tensor shapes"
]
@impl Exception
def message(%__MODULE__{} = error) do
parts = ["Shape mismatch in #{error.operation}"]
parts =
if error.shape_a do
parts ++ [" Shape A: #{inspect(error.shape_a)}"]
else
parts
end
parts =
if error.shape_b do
parts ++ [" Shape B: #{inspect(error.shape_b)}"]
else
parts
end
parts =
if error.expected do
parts ++ [" Expected: #{error.expected}"]
else
parts
end
parts =
if error.got do
parts ++ [" Got: #{error.got}"]
else
parts
end
parts = parts ++ ["", error.message, "", "Suggestion: #{error.suggestion}"]
parts
|> Enum.join("\n")
|> String.trim()
end
@doc """
Creates a ShapeMismatchError error from context.
"""
@spec new(atom(), keyword()) :: t()
def new(operation, opts \\ []) do
shape_a = Keyword.get(opts, :shape_a)
shape_b = Keyword.get(opts, :shape_b)
suggestion =
Keyword.get(opts, :suggestion) ||
generate_suggestion(operation, shape_a, shape_b)
%__MODULE__{
operation: operation,
shape_a: shape_a,
shape_b: shape_b,
expected: Keyword.get(opts, :expected),
got: Keyword.get(opts, :got),
message: Keyword.get(opts, :message, "Shapes are incompatible for #{operation}"),
suggestion: suggestion,
python_traceback: Keyword.get(opts, :python_traceback)
}
end
@doc """
Generates a suggestion based on the operation and shapes.
"""
@spec generate_suggestion(atom(), [non_neg_integer()] | nil, [non_neg_integer()] | nil) ::
String.t()
def generate_suggestion(:matmul, shape_a, shape_b)
when is_list(shape_a) and is_list(shape_b) do
a_cols = List.last(shape_a)
b_rows = List.first(shape_b)
if a_cols != b_rows do
"For matrix multiplication, A columns (#{a_cols}) must equal B rows (#{b_rows}). " <>
"Try: tensor.transpose(dim0, dim1) if B needs transposing"
else
"Check that tensor shapes are compatible for matrix multiplication."
end
end
def generate_suggestion(_operation, shape_a, shape_b)
when is_list(shape_a) and is_list(shape_b) do
if length(shape_a) != length(shape_b) do
"Tensors have different number of dimensions (#{length(shape_a)} vs #{length(shape_b)}). " <>
"Use unsqueeze/squeeze to adjust dimensions."
else
mismatched = find_mismatched_dim(shape_a, shape_b)
if mismatched do
"Shapes differ at dimension #{mismatched}. Check broadcasting rules or reshape tensors."
else
"Verify tensor shapes are compatible for this operation."
end
end
end
def generate_suggestion(_operation, _shape_a, _shape_b) do
"Verify tensor shapes are compatible for this operation."
end
defp find_mismatched_dim(shape_a, shape_b) do
shape_a
|> Enum.zip(shape_b)
|> Enum.with_index()
|> Enum.find_value(fn
{{x, y}, idx} when x != y and x != 1 and y != 1 -> idx
_ -> nil
end)
end
end
defmodule SnakeBridge.Generator.TypeMapper do
@moduledoc """
Maps Python type annotations to Elixir typespec AST.
This module converts Python type dictionaries (as produced by the introspection
script) into Elixir typespec AST using `quote`. The AST can then be used to
generate `@spec` declarations in generated modules.
## Type Mappings
| Python Type | Elixir Type |
|------------|-------------|
| `int` | `integer()` |
| `float` | `float()` |
| `str` | `String.t()` |
| `bool` | `boolean()` |
| `bytes` | `binary()` |
| `None` | `nil` |
| `list[T]` | `list(T)` |
| `dict[K, V]` | `map(K, V)` |
| `tuple[T1, T2, ...]` | `{T1, T2, ...}` |
| `set[T]` | `MapSet.t(T)` |
| `Optional[T]` | `T \\| nil` |
| `Union[T1, T2, ...]` | `T1 \\| T2 \\| ...` |
| `ClassName` | `ClassName.t()` |
| `Any` | `any()` |
## Examples
iex> TypeMapper.to_spec(%{"type" => "int"})
{:integer, [], []}
iex> TypeMapper.to_spec(%{"type" => "list", "element_type" => %{"type" => "str"}})
{{:., [], [{:__aliases__, [alias: false], [:String]}, :t]}, [], []}
|> Macro.to_string()
"list(String.t())"
"""
@doc """
Converts a Python type dictionary to an Elixir typespec AST.
## Parameters
* `python_type` - A map representing a Python type annotation
## Returns
An AST node (quoted expression) representing the equivalent Elixir typespec.
## Examples
iex> python_type = %{"type" => "int"}
iex> ast = SnakeBridge.Generator.TypeMapper.to_spec(python_type)
iex> Macro.to_string(ast)
"integer()"
iex> python_type = %{"type" => "list", "element_type" => %{"type" => "int"}}
iex> ast = SnakeBridge.Generator.TypeMapper.to_spec(python_type)
iex> Macro.to_string(ast)
"list(integer())"
"""
@spec to_spec(map() | nil) :: Macro.t()
def to_spec(nil), do: quote(do: term())
def to_spec(%{} = python_type) when map_size(python_type) == 0, do: quote(do: term())
# Primitive types
def to_spec(%{"type" => "int"}), do: quote(do: integer())
def to_spec(%{"type" => "float"}), do: quote(do: float())
def to_spec(%{"type" => "str"}), do: quote(do: String.t())
def to_spec(%{"type" => "string"}), do: quote(do: String.t())
def to_spec(%{"type" => "bool"}), do: quote(do: boolean())
def to_spec(%{"type" => "boolean"}), do: quote(do: boolean())
def to_spec(%{"type" => "bytes"}), do: quote(do: binary())
def to_spec(%{"type" => "bytearray"}), do: quote(do: binary())
def to_spec(%{"type" => "none"}), do: quote(do: nil)
def to_spec(%{"type" => "any"}), do: quote(do: term())
# Complex types - delegate to specialized mappers
def to_spec(%{"type" => "list"} = python_type), do: map_list_type(python_type)
def to_spec(%{"type" => "dict"} = python_type), do: map_dict_type(python_type)
def to_spec(%{"type" => "tuple"} = python_type), do: map_tuple_type(python_type)
def to_spec(%{"type" => "set"} = python_type), do: map_set_type(python_type)
def to_spec(%{"type" => "frozenset"} = python_type), do: map_set_type(python_type)
def to_spec(%{"type" => "optional"} = python_type), do: map_optional_type(python_type)
def to_spec(%{"type" => "union"} = python_type), do: map_union_type(python_type)
def to_spec(%{"type" => "class"} = python_type), do: map_class_type(python_type)
# ML-specific types (NumPy, PyTorch, Pandas)
def to_spec(%{"type" => "numpy.ndarray"}), do: quote(do: Numpy.NDArray.t())
def to_spec(%{"type" => "numpy.dtype"}), do: quote(do: Numpy.DType.t())
def to_spec(%{"type" => "torch.tensor"}), do: quote(do: Torch.Tensor.t())
def to_spec(%{"type" => "torch.Tensor"}), do: quote(do: Torch.Tensor.t())
def to_spec(%{"type" => "torch.dtype"}), do: quote(do: Torch.DType.t())
def to_spec(%{"type" => "pandas.dataframe"}), do: quote(do: Pandas.DataFrame.t())
def to_spec(%{"type" => "pandas.DataFrame"}), do: quote(do: Pandas.DataFrame.t())
def to_spec(%{"type" => "pandas.series"}), do: quote(do: Pandas.Series.t())
def to_spec(%{"type" => "pandas.Series"}), do: quote(do: Pandas.Series.t())
# Python integer alias
def to_spec(%{"type" => "integer"}), do: quote(do: integer())
# Fallback for unknown types
def to_spec(%{"type" => _}), do: quote(do: term())
def to_spec(_), do: quote(do: term())
# Private Functions
@spec map_list_type(map()) :: Macro.t()
defp map_list_type(%{"element_type" => element_type}) do
element_spec = to_spec(element_type)
quote(do: list(unquote(element_spec)))
end
defp map_list_type(_), do: quote(do: list(term()))
@spec map_dict_type(map()) :: Macro.t()
defp map_dict_type(%{"key_type" => key_type, "value_type" => value_type}) do
key_spec = to_spec(key_type)
value_spec = to_spec(value_type)
quote(do: %{optional(unquote(key_spec)) => unquote(value_spec)})
end
defp map_dict_type(_), do: quote(do: %{optional(term()) => term()})
@spec map_tuple_type(map()) :: Macro.t()
defp map_tuple_type(%{"element_types" => element_types}) when is_list(element_types) do
case element_types do
[] ->
{:{}, [], []}
types ->
element_specs = Enum.map(types, &to_spec/1)
{:{}, [], element_specs}
end
end
defp map_tuple_type(_), do: quote(do: tuple())
@spec map_set_type(map()) :: Macro.t()
defp map_set_type(%{"element_type" => element_type}) do
element_spec = to_spec(element_type)
quote(do: MapSet.t(unquote(element_spec)))
end
defp map_set_type(_), do: quote(do: MapSet.t(term()))
@spec map_optional_type(map()) :: Macro.t()
defp map_optional_type(%{"inner_type" => inner_type}) do
inner_spec = to_spec(inner_type)
quote(do: unquote(inner_spec) | nil)
end
defp map_optional_type(_), do: quote(do: term() | nil)
@spec map_union_type(map()) :: Macro.t()
defp map_union_type(%{"types" => types}) when is_list(types) and length(types) > 0 do
type_specs = Enum.map(types, &to_spec/1)
# Build union type using |
Enum.reduce(type_specs, fn spec, acc ->
quote(do: unquote(acc) | unquote(spec))
end)
end
defp map_union_type(_), do: quote(do: term())
@spec map_class_type(map()) :: Macro.t()
defp map_class_type(%{"name" => name, "module" => module})
when is_binary(name) and is_binary(module) do
module_parts =
module
|> String.split(".")
|> Enum.map(&Macro.camelize/1)
|> Kernel.++([name])
module_alias = {:__aliases__, [alias: false], Enum.map(module_parts, &String.to_atom/1)}
{{:., [], [module_alias, :t]}, [], []}
end
defp map_class_type(%{"name" => name}) when is_binary(name) do
module_alias = {:__aliases__, [alias: false], [String.to_atom(name)]}
{{:., [], [module_alias, :t]}, [], []}
end
defp map_class_type(_), do: quote(do: term())
end
defmodule SnakeBridge.Lock.Verifier do
@moduledoc """
Verifies hardware and environment compatibility between the lock file and current system.
The verifier compares the hardware identity in the lock file against the current
system's capabilities to detect potential compatibility issues before runtime.
## Verification Levels
- `:ok` - Full compatibility, no issues detected
- `{:warning, warnings}` - Minor differences that may work but could cause issues
- `{:error, errors}` - Incompatible environment that will likely fail
## Examples
# Verify lock file compatibility
lock = SnakeBridge.Lock.load()
case SnakeBridge.Lock.Verifier.verify(lock) do
:ok ->
IO.puts("Environment compatible")
{:warning, warnings} ->
Enum.each(warnings, &IO.warn/1)
{:error, errors} ->
raise SnakeBridge.EnvironmentError, message: Enum.join(errors, "; ")
end
"""
@type verification_result :: :ok | {:warning, [String.t()]} | {:error, [String.t()]}
@doc """
Verifies the lock file against the current hardware environment.
Returns `:ok` if compatible, `{:warning, warnings}` for minor issues,
or `{:error, errors}` for critical incompatibilities.
"""
@spec verify(map() | nil) :: verification_result()
def verify(nil) do
start_time = System.monotonic_time()
SnakeBridge.Telemetry.lock_verify(start_time, :error, ["No lock file provided"])
{:error, ["No lock file provided"]}
end
def verify(lock) when is_map(lock) do
start_time = System.monotonic_time()
errors = []
warnings = []
current = hardware_module().identity()
lock_env = Map.get(lock, "environment", %{})
lock_hardware = Map.get(lock_env, "hardware", %{})
lock_platform = Map.get(lock_env, "platform", %{})
compatibility = Map.get(lock, "compatibility", %{})
# Check platform
{platform_errors, platform_warnings} = verify_platform(lock_platform, current)
errors = errors ++ platform_errors
warnings = warnings ++ platform_warnings
# Check accelerator/CUDA
{accel_errors, accel_warnings} = verify_accelerator(lock_hardware, current, compatibility)
errors = errors ++ accel_errors
warnings = warnings ++ accel_warnings
# Check CPU features if required
{feature_errors, feature_warnings} = verify_cpu_features(lock_hardware, current)
errors = errors ++ feature_errors
warnings = warnings ++ feature_warnings
result =
cond do
errors != [] -> {:error, errors}
warnings != [] -> {:warning, warnings}
true -> :ok
end
case result do
:ok -> SnakeBridge.Telemetry.lock_verify(start_time, :ok, [])
{:warning, warn} -> SnakeBridge.Telemetry.lock_verify(start_time, :warning, warn)
{:error, errs} -> SnakeBridge.Telemetry.lock_verify(start_time, :error, errs)
end
result
end
@doc """
Verifies the lock file and raises on error.
Returns `:ok` on success or raises `SnakeBridge.EnvironmentError`.
Warnings are logged but do not raise.
"""
@spec verify!(map() | nil) :: :ok
def verify!(lock) do
case verify(lock) do
:ok ->
:ok
{:warning, warnings} ->
Enum.each(warnings, &Mix.shell().info("Warning: #{&1}"))
:ok
{:error, errors} ->
raise SnakeBridge.EnvironmentError,
message: "Lock file incompatible: #{Enum.join(errors, "; ")}"
end
end
# Private functions
defp verify_platform(lock_platform, _current) when map_size(lock_platform) == 0 do
{[], []}
end
defp verify_platform(lock_platform, current) do
errors = []
warnings = []
current_platform = current["platform"] || ""
[current_os, current_arch] = parse_platform(current_platform)
lock_os = Map.get(lock_platform, "os", "")
lock_arch = Map.get(lock_platform, "arch", "")
errors =
if lock_os != "" and lock_os != current_os do
["Platform mismatch: lock requires #{lock_os}, current is #{current_os}" | errors]
else
errors
end
errors =
if lock_arch != "" and lock_arch != current_arch do
["Architecture mismatch: lock requires #{lock_arch}, current is #{current_arch}" | errors]
else
errors
end
{errors, warnings}
end
defp verify_accelerator(lock_hardware, _current, _compatibility)
when map_size(lock_hardware) == 0 do
{[], []}
end
defp verify_accelerator(lock_hardware, current, _compatibility) do
lock_accelerator = Map.get(lock_hardware, "accelerator", "cpu")
current_accelerator = current["accelerator"] || "cpu"
current_caps = hardware_module().capabilities()
lock_cuda_version = Map.get(lock_hardware, "cuda_version")
current_cuda_version = current_caps.cuda_version
check_accelerator_compatibility(
lock_accelerator,
current_accelerator,
current_caps,
lock_cuda_version,
current_cuda_version
)
end
defp check_accelerator_compatibility(
"cuda",
_current_accel,
%{cuda: false},
_lock_ver,
_cur_ver
) do
{["Lock requires CUDA but no CUDA available on current system"], []}
end
defp check_accelerator_compatibility("mps", _current_accel, %{mps: false}, _lock_ver, _cur_ver) do
{["Lock requires MPS but MPS not available (requires macOS with Apple Silicon)"], []}
end
defp check_accelerator_compatibility("cuda", _current_accel, %{cuda: true}, lock_ver, cur_ver) do
check_cuda_version_compatibility(lock_ver, cur_ver)
end
defp check_accelerator_compatibility("cuda", "cpu", _caps, _lock_ver, _cur_ver) do
{[], ["Lock was built with CUDA, falling back to CPU"]}
end
defp check_accelerator_compatibility(_lock_accel, _current_accel, _caps, _lock_ver, _cur_ver) do
{[], []}
end
defp check_cuda_version_compatibility(lock_version, current_version) do
lock_major = major_version(lock_version)
current_major = major_version(current_version)
cond do
lock_major != current_major ->
{[], ["CUDA version mismatch: lock has #{lock_version}, current has #{current_version}"]}
lock_version != current_version ->
{[], ["CUDA version differs: lock has #{lock_version}, current has #{current_version}"]}
true ->
{[], []}
end
end
defp verify_cpu_features(lock_hardware, current) do
lock_features_list = Map.get(lock_hardware, "cpu_features", [])
current_features_list = Map.get(current, "cpu_features", [])
critical_list = ["avx512f"]
lock_features = MapSet.new(lock_features_list)
current_features = MapSet.new(current_features_list)
critical_features = MapSet.new(critical_list)
missing_critical =
lock_features
|> MapSet.intersection(critical_features)
|> MapSet.difference(current_features)
|> MapSet.to_list()
missing_optional =
lock_features
|> MapSet.difference(critical_features)
|> MapSet.difference(current_features)
|> MapSet.to_list()
errors =
if missing_critical != [] do
["Missing critical CPU features: #{Enum.join(missing_critical, ", ")}"]
else
[]
end
warnings =
if missing_optional != [] do
["Missing optional CPU features: #{Enum.join(missing_optional, ", ")}"]
else
[]
end
{errors, warnings}
end
defp parse_platform(platform_string) when is_binary(platform_string) do
case String.split(platform_string, "-", parts: 2) do
[os, arch] -> [os, arch]
[os] -> [os, "unknown"]
_ -> ["unknown", "unknown"]
end
end
defp parse_platform(_), do: ["unknown", "unknown"]
defp major_version(nil), do: nil
defp major_version(version) when is_binary(version) do
case String.split(version, ".") do
[major | _] -> major
_ -> version
end
end
defp hardware_module do
Application.get_env(:snakebridge, :hardware_module, Snakepit.Hardware)
end
end
defmodule SnakeBridge.PythonRunner.System do
@moduledoc false
@behaviour SnakeBridge.PythonRunner
@impl SnakeBridge.PythonRunner
def run(script, args, opts \\ []) when is_binary(script) and is_list(args) do
with {:ok, python, _meta} <- Snakepit.PythonRuntime.resolve_executable() do
env = build_env(opts)
cmd_opts = Keyword.merge([stderr_to_stdout: true, env: env], Keyword.drop(opts, [:env]))
case System.cmd(python, ["-c", script | args], cmd_opts) do
{output, 0} -> {:ok, output}
{output, status} -> {:error, {:python_exit, status, output}}
end
end
end
defp build_env(opts) do
runtime_env = Snakepit.PythonRuntime.runtime_env()
extra_env =
Snakepit.PythonRuntime.config()
|> Map.get(:extra_env, %{})
|> Enum.to_list()
user_env =
opts
|> Keyword.get(:env, %{})
|> Enum.to_list()
runtime_env ++ extra_env ++ user_env
end
end
defmodule SnakeBridge.Telemetry.Handlers.Logger do
@moduledoc """
Logs SnakeBridge telemetry events.
This handler logs compilation events at appropriate log levels:
- Compile stop: `:info`
- Compile exception: `:error`
- Introspect/Generate: `:debug`
## Usage
# In your application startup
SnakeBridge.Telemetry.Handlers.Logger.attach()
"""
require Logger
@handler_id "snakebridge-logger"
@events [
[:snakebridge, :compile, :stop],
[:snakebridge, :compile, :exception],
[:snakebridge, :compile, :introspect, :stop],
[:snakebridge, :compile, :generate, :stop]
]
@doc """
Attaches the logger handler to telemetry events.
Returns `:ok` on success or `{:error, :already_exists}` if already attached.
"""
@spec attach() :: :ok | {:error, :already_exists}
def attach do
:telemetry.attach_many(
@handler_id,
@events,
&handle_event/4,
%{}
)
end
@doc """
Detaches the logger handler.
"""
@spec detach() :: :ok | {:error, :not_found}
def detach do
:telemetry.detach(@handler_id)
end
@doc false
def handle_event([:snakebridge, :compile, :stop], measurements, metadata, _config) do
duration_ms = System.convert_time_unit(measurements.duration, :native, :millisecond)
libraries = metadata.details[:libraries] || []
Logger.info(
"SnakeBridge compiled #{measurements.symbols_generated} symbols " <>
"in #{duration_ms}ms (#{length(libraries)} libraries)"
)
end
def handle_event([:snakebridge, :compile, :exception], _measurements, metadata, _config) do
reason = metadata.details[:reason]
Logger.error("SnakeBridge compilation failed: #{inspect(reason)}")
end
def handle_event([:snakebridge, :compile, :introspect, :stop], measurements, metadata, _config) do
duration_ms = System.convert_time_unit(measurements.duration, :native, :millisecond)
python_time = metadata.details[:python_time] || 0
python_ms = System.convert_time_unit(python_time, :native, :millisecond)
Logger.debug(
"Introspected #{measurements.symbols_introspected} symbols from #{metadata.library} " <>
"in #{duration_ms}ms (Python: #{python_ms}ms, cache hits: #{measurements.cache_hits})"
)
end
def handle_event([:snakebridge, :compile, :generate, :stop], measurements, metadata, _config) do
duration_ms = System.convert_time_unit(measurements.duration, :native, :millisecond)
Logger.debug(
"Generated #{measurements.functions_generated} functions for #{metadata.library} " <>
"in #{duration_ms}ms (#{measurements.bytes_written} bytes)"
)
end
end
defmodule SnakeBridge.Telemetry.Handlers.Metrics do
@moduledoc """
Metric definitions for SnakeBridge telemetry.
This module provides metric definitions compatible with TelemetryMetrics
and reporters like TelemetryMetricsPrometheus.
## Usage
# In your application with TelemetryMetricsPrometheus
TelemetryMetricsPrometheus.Core.attach(
SnakeBridge.Telemetry.Handlers.Metrics.metrics()
)
## Metrics
### Compilation
- `snakebridge.compile.duration` - Distribution of compilation times
- `snakebridge.compile.symbols_generated` - Sum of symbols generated
- `snakebridge.compile.total` - Counter of compilations
### Scanning
- `snakebridge.scan.duration` - Distribution of scan times
- `snakebridge.scan.files_scanned` - Sum of files scanned
- `snakebridge.scan.symbols_found` - Sum of symbols found
### Introspection
- `snakebridge.introspect.duration` - Distribution of introspection times
- `snakebridge.introspect.symbols_introspected` - Sum of symbols introspected
- `snakebridge.introspect.cache_hits` - Sum of cache hits
### Generation
- `snakebridge.generate.duration` - Distribution of generation times
- `snakebridge.generate.bytes_written` - Sum of bytes written
### Documentation
- `snakebridge.docs.fetch.duration` - Distribution of doc fetch times
- `snakebridge.docs.fetch.total` - Counter of doc fetches
"""
@doc """
Returns a list of Telemetry.Metrics definitions.
These can be used with any TelemetryMetrics-compatible reporter.
"""
@spec metrics() :: [struct()]
def metrics do
import Telemetry.Metrics
[
# Compilation metrics
distribution("snakebridge.compile.duration",
event_name: [:snakebridge, :compile, :stop],
measurement: :duration,
unit: {:native, :millisecond},
reporter_options: [buckets: [100, 500, 1000, 5000, 10_000]]
),
sum("snakebridge.compile.symbols_generated",
event_name: [:snakebridge, :compile, :stop],
measurement: :symbols_generated
),
counter("snakebridge.compile.total",
event_name: [:snakebridge, :compile, :stop]
),
# Scan metrics
distribution("snakebridge.scan.duration",
event_name: [:snakebridge, :compile, :scan, :stop],
measurement: :duration,
unit: {:native, :millisecond}
),
sum("snakebridge.scan.files_scanned",
event_name: [:snakebridge, :compile, :scan, :stop],
measurement: :files_scanned
),
sum("snakebridge.scan.symbols_found",
event_name: [:snakebridge, :compile, :scan, :stop],
measurement: :symbols_found
),
# Introspection metrics
distribution("snakebridge.introspect.duration",
event_name: [:snakebridge, :compile, :introspect, :stop],
measurement: :duration,
tags: [:library],
unit: {:native, :millisecond}
),
sum("snakebridge.introspect.symbols_introspected",
event_name: [:snakebridge, :compile, :introspect, :stop],
measurement: :symbols_introspected,
tags: [:library]
),
sum("snakebridge.introspect.cache_hits",
event_name: [:snakebridge, :compile, :introspect, :stop],
measurement: :cache_hits,
tags: [:library]
),
# Generation metrics
distribution("snakebridge.generate.duration",
event_name: [:snakebridge, :compile, :generate, :stop],
measurement: :duration,
tags: [:library],
unit: {:native, :millisecond}
),
sum("snakebridge.generate.bytes_written",
event_name: [:snakebridge, :compile, :generate, :stop],
measurement: :bytes_written,
tags: [:library]
),
# Documentation metrics
distribution("snakebridge.docs.fetch.duration",
event_name: [:snakebridge, :docs, :fetch],
measurement: :duration,
tags: [:source],
unit: {:native, :millisecond}
),
counter("snakebridge.docs.fetch.total",
event_name: [:snakebridge, :docs, :fetch],
tags: [:source]
)
]
end
end
defmodule SnakeBridge.Telemetry.RuntimeForwarder do
@moduledoc """
Enriches Snakepit runtime telemetry with SnakeBridge context.
This module listens to Snakepit's call events and re-emits them under
the `:snakebridge` namespace with additional context like the SnakeBridge
version and library information.
## Events
Original Snakepit events:
- `[:snakepit, :python, :call, :start]`
- `[:snakepit, :python, :call, :stop]`
- `[:snakepit, :python, :call, :exception]`
Are forwarded as:
- `[:snakebridge, :runtime, :call, :start]`
- `[:snakebridge, :runtime, :call, :stop]`
- `[:snakebridge, :runtime, :call, :exception]`
With added metadata:
- `snakebridge_library` - The library name from the original event
- `snakebridge_version` - The current SnakeBridge version
## Usage
# In your application startup
SnakeBridge.Telemetry.RuntimeForwarder.attach()
"""
@handler_id "snakebridge-runtime-enricher"
@events [
[:snakepit, :python, :call, :start],
[:snakepit, :python, :call, :stop],
[:snakepit, :python, :call, :exception]
]
@doc """
Attaches the runtime forwarder to Snakepit events.
Returns `:ok` on success or `{:error, :already_exists}` if already attached.
"""
@spec attach() :: :ok | {:error, :already_exists}
def attach do
:telemetry.attach_many(
@handler_id,
@events,
&handle_event/4,
%{}
)
end
@doc """
Detaches the runtime forwarder.
"""
@spec detach() :: :ok | {:error, :not_found}
def detach do
:telemetry.detach(@handler_id)
end
@doc false
def handle_event([:snakepit, :python, :call, :start], measurements, metadata, _config) do
enriched = enrich_metadata(metadata)
:telemetry.execute(
[:snakebridge, :runtime, :call, :start],
measurements,
enriched
)
end
def handle_event([:snakepit, :python, :call, :stop], measurements, metadata, _config) do
enriched = enrich_metadata(metadata)
:telemetry.execute(
[:snakebridge, :runtime, :call, :stop],
measurements,
enriched
)
end
def handle_event([:snakepit, :python, :call, :exception], measurements, metadata, _config) do
enriched = enrich_metadata(metadata)
:telemetry.execute(
[:snakebridge, :runtime, :call, :exception],
measurements,
enriched
)
end
defp enrich_metadata(metadata) do
library =
Map.get(metadata, :library) || Map.get(metadata, :snakebridge_library) ||
Map.get(metadata, :python_module)
function = Map.get(metadata, :function) || Map.get(metadata, :name)
call_type = Map.get(metadata, :call_type) || Map.get(metadata, :type)
metadata
|> Map.put(:library, library || "unknown")
|> Map.put(:function, function || "unknown")
|> Map.put(:call_type, call_type || "unknown")
|> Map.put(:snakebridge_library, library || "unknown")
|> Map.put(:snakebridge_version, version())
end
defp version do
Application.spec(:snakebridge, :vsn) |> to_string()
end
end
defmodule SnakeBridge.Types.Decoder do
@moduledoc """
Decodes JSON-compatible data from Python into Elixir data structures.
Handles lossless decoding of tagged representations produced by the Python
side or by `SnakeBridge.Types.Encoder`. Recognizes special `__type__` markers
to reconstruct Elixir-specific types. Atom decoding is allowlist-based
(configure via `:snakebridge, :atom_allowlist`).
## Supported Tagged Types
- `{"__type__": "atom", "value": "ok"}` → `:ok` (allowlisted only)
- `{"__type__": "tuple", "elements": [...]}` → Elixir tuple
- `{"__type__": "set", "elements": [...]}` → MapSet
- `{"__type__": "frozenset", "elements": [...]}` → MapSet
- `{"__type__": "bytes", "data": ""}` → binary
- `{"__type__": "datetime", "value": ""}` → DateTime
- `{"__type__": "date", "value": ""}` → Date
- `{"__type__": "time", "value": ""}` → Time
- `{"__type__": "special_float", "value": "infinity"}` → `:infinity`
- `{"__type__": "special_float", "value": "neg_infinity"}` → `:neg_infinity`
- `{"__type__": "special_float", "value": "nan"}` → `:nan`
- `{"__type__": "ref", ...}` → `SnakeBridge.Ref`
- `{"__type__": "stream_ref", ...}` → `SnakeBridge.StreamRef`
## Direct JSON Types
- `null` → `nil`
- Booleans → `true`/`false`
- Numbers → integers or floats
- Strings → strings
- Arrays → lists (recursively decoded)
- Objects → maps with string keys (recursively decoded)
## Examples
iex> SnakeBridge.Types.Decoder.decode(%{"__type__" => "tuple", "elements" => [1, 2, 3]})
{1, 2, 3}
iex> SnakeBridge.Types.Decoder.decode(%{"__type__" => "set", "elements" => [1, 2, 3]})
#MapSet<[1, 2, 3]>
iex> SnakeBridge.Types.Decoder.decode(%{"a" => 1, "b" => 2})
%{"a" => 1, "b" => 2}
"""
@doc """
Decodes a JSON-compatible value into an Elixir data structure.
Recognizes and handles tagged types from the Python encoder.
## Examples
iex> decode(42)
42
iex> decode([1, 2, 3])
[1, 2, 3]
iex> decode(%{
...> "__type__" => "tuple",
...> "elements" => [%{"__type__" => "atom", "value" => "ok"}, "result"]
...> })
{:ok, "result"}
"""
@spec decode(term()) :: term()
def decode(nil), do: nil
def decode(true), do: true
def decode(false), do: false
def decode(num) when is_number(num), do: num
def decode(str) when is_binary(str), do: str
# Lists - recursively decode elements
def decode(list) when is_list(list) do
Enum.map(list, &decode/1)
end
def decode(%{"__type__" => "stream_ref"} = map) do
SnakeBridge.StreamRef.from_wire_format(map)
end
def decode(%{"__type__" => "ref"} = ref) do
required = ["id", "session_id"]
missing = Enum.filter(required, &(not Map.has_key?(ref, &1)))
if missing != [] do
raise ArgumentError, "Invalid ref: missing fields #{inspect(missing)}"
end
ref_struct = SnakeBridge.Ref.from_wire_format(ref)
maybe_register_ref(ref_struct)
ref_struct
end
# Maps with __type__ markers - decode based on type
def decode(%{"__type__" => "atom"} = map) do
value = Map.get(map, "value")
if is_binary(value) and atom_allowed?(value) do
String.to_atom(value)
else
value
end
end
def decode(%{"__type__" => "tuple"} = map) do
map
|> list_field()
|> Enum.map(&decode/1)
|> List.to_tuple()
end
def decode(%{"__type__" => "set"} = map) do
map
|> list_field()
|> Enum.map(&decode/1)
|> MapSet.new()
end
def decode(%{"__type__" => "frozenset"} = map) do
map
|> list_field()
|> Enum.map(&decode/1)
|> MapSet.new()
end
def decode(%{"__type__" => "bytes"} = map) do
data = Map.get(map, "data") || Map.get(map, "value")
if is_binary(data) do
case Base.decode64(data) do
{:ok, binary} -> binary
:error -> data
end
else
data
end
end
def decode(%{"__type__" => "datetime", "value" => value}) when is_binary(value) do
case DateTime.from_iso8601(value) do
{:ok, dt, _offset} -> dt
{:error, _} -> value
end
end
def decode(%{"__type__" => "date", "value" => value}) when is_binary(value) do
case Date.from_iso8601(value) do
{:ok, date} -> date
{:error, _} -> value
end
end
def decode(%{"__type__" => "time", "value" => value}) when is_binary(value) do
case Time.from_iso8601(value) do
{:ok, time} -> time
{:error, _} -> value
end
end
def decode(%{"__type__" => "special_float", "value" => "infinity"}), do: :infinity
def decode(%{"__type__" => "special_float", "value" => "neg_infinity"}), do: :neg_infinity
def decode(%{"__type__" => "special_float", "value" => "nan"}), do: :nan
def decode(%{"__type__" => "infinity"}), do: :infinity
def decode(%{"__type__" => "neg_infinity"}), do: :neg_infinity
def decode(%{"__type__" => "nan"}), do: :nan
def decode(%{"__type__" => "complex", "real" => real, "imag" => imag}) do
%{real: real, imag: imag}
end
# Tagged dict - maps with non-string keys
def decode(%{"__type__" => "dict", "pairs" => pairs}) when is_list(pairs) do
pairs
|> Enum.map(fn
[key, value] ->
{decode(key), decode(value)}
pair when is_list(pair) and length(pair) == 2 ->
[key, value] = pair
{decode(key), decode(value)}
end)
|> Map.new()
end
# Tagged dict with schema version
def decode(%{"__type__" => "dict", "__schema__" => _schema, "pairs" => pairs}) do
decode(%{"__type__" => "dict", "pairs" => pairs})
end
# Regular maps - recursively decode values
def decode(%{} = map) do
Map.new(map, fn {key, value} ->
{key, decode(value)}
end)
end
# Anything else passes through unchanged
def decode(other), do: other
defp list_field(map) do
case Map.get(map, "elements") do
nil -> Map.get(map, "value", [])
elements -> elements
end
|> List.wrap()
end
defp atom_allowed?(value) when is_binary(value) do
case atom_allowlist() do
:all -> true
allowlist -> value in allowlist
end
end
defp atom_allowed?(_), do: false
defp atom_allowlist do
case Application.get_env(:snakebridge, :atom_allowlist, ["ok", "error"]) do
:all -> :all
list -> Enum.map(List.wrap(list), &to_string/1)
end
end
defp maybe_register_ref(ref) do
session_id = Map.get(ref, "session_id") || Map.get(ref, :session_id)
if is_binary(session_id) and Process.whereis(SnakeBridge.SessionManager) do
case SnakeBridge.SessionManager.register_ref(session_id, ref) do
:ok -> :ok
{:error, _reason} -> :ok
end
end
:ok
end
end
defmodule SnakeBridge.Types.Encoder do
@moduledoc """
Encodes Elixir data structures into JSON-compatible formats for Python interop.
Handles lossless encoding of Elixir types that don't have direct JSON equivalents
using tagged representations. Tagged values include a `__schema__` marker for
the current wire schema version. Atom round-trips depend on the decoder
allowlist.
## Supported Types
### Direct JSON Types
- `nil` → `null`
- Booleans → `true`/`false`
- Integers → numbers
- Floats → numbers
- Strings (UTF-8) → strings
- Lists → arrays
- Maps with string keys → objects
### Tagged Types
- Atoms → `{"__type__": "atom", "value": "ok"}`
- Tuples → `{"__type__": "tuple", "elements": [...]}`
- MapSets → `{"__type__": "set", "elements": [...]}`
- Binaries (non-UTF-8) → `{"__type__": "bytes", "data": ""}`
- `SnakeBridge.Bytes` → `{"__type__": "bytes", "data": ""}` (always bytes)
- DateTime → `{"__type__": "datetime", "value": ""}`
- Date → `{"__type__": "date", "value": ""}`
- Time → `{"__type__": "time", "value": ""}`
- Special floats → `{"__type__": "special_float", "value": "infinity"|"neg_infinity"|"nan"}`
- Maps with string/atom keys → plain objects (keys converted to strings)
- Maps with non-string keys → `{"__type__": "dict", "pairs": [[key, val], ...]}`
## Unsupported Types
The following types cannot be serialized and will raise `SnakeBridge.SerializationError`:
- PIDs, ports, references
- Custom structs without explicit encoder support
## Examples
iex> SnakeBridge.Types.Encoder.encode(%{a: 1, b: 2})
%{"a" => 1, "b" => 2}
iex> SnakeBridge.Types.Encoder.encode({:ok, "result"})
%{
"__type__" => "tuple",
"__schema__" => 1,
"elements" => [%{"__type__" => "atom", "__schema__" => 1, "value" => "ok"}, "result"]
}
iex> SnakeBridge.Types.Encoder.encode(MapSet.new([1, 2, 3]))
%{"__type__" => "set", "__schema__" => 1, "elements" => [1, 2, 3]}
iex> SnakeBridge.Types.Encoder.encode(%{1 => "one", 2 => "two"})
%{"__type__" => "dict", "__schema__" => 1, "pairs" => [[1, "one"], [2, "two"]]}
"""
@doc """
Encodes an Elixir value into a JSON-compatible structure.
## Examples
iex> encode(42)
42
iex> encode(:ok)
%{"__type__" => "atom", "__schema__" => 1, "value" => "ok"}
iex> encode({1, 2, 3})
%{"__type__" => "tuple", "__schema__" => 1, "elements" => [1, 2, 3]}
## Raises
- `SnakeBridge.SerializationError` for unsupported types (PIDs, ports, refs, unknown structs)
"""
@spec encode(term()) :: term()
def encode(nil), do: nil
def encode(true), do: true
def encode(false), do: false
# Special float atoms
def encode(:infinity), do: tagged("special_float", %{"value" => "infinity"})
def encode(:neg_infinity), do: tagged("special_float", %{"value" => "neg_infinity"})
def encode(:nan), do: tagged("special_float", %{"value" => "nan"})
# Regular atoms are tagged for lossless interop
def encode(atom) when is_atom(atom) do
tagged("atom", %{"value" => Atom.to_string(atom)})
end
# Numbers
def encode(num) when is_integer(num), do: num
def encode(num) when is_float(num), do: num
# Explicit bytes wrapper - MUST come before generic binary clause
def encode(%SnakeBridge.Bytes{data: data}) when is_binary(data) do
tagged("bytes", %{"data" => Base.encode64(data)})
end
# Strings and binaries
def encode(binary) when is_binary(binary) do
if String.valid?(binary) do
binary
else
# Non-UTF-8 binary - encode as base64
tagged("bytes", %{"data" => Base.encode64(binary)})
end
end
# Lists
def encode(list) when is_list(list) do
Enum.map(list, &encode/1)
end
# Tuples
def encode(tuple) when is_tuple(tuple) do
tagged("tuple", %{"elements" => tuple |> Tuple.to_list() |> Enum.map(&encode/1)})
end
# MapSets
def encode(%MapSet{} = mapset) do
tagged("set", %{"elements" => mapset |> MapSet.to_list() |> Enum.map(&encode/1)})
end
# DateTime
def encode(%DateTime{} = dt) do
tagged("datetime", %{"value" => DateTime.to_iso8601(dt)})
end
# Date
def encode(%Date{} = date) do
tagged("date", %{"value" => Date.to_iso8601(date)})
end
# Time
def encode(%Time{} = time) do
tagged("time", %{"value" => Time.to_iso8601(time)})
end
# Snakepit PyRef - normalize to ref wire shape
def encode(%{__struct__: Snakepit.PyRef} = ref) do
ref
|> Map.from_struct()
|> normalize_pyref_map()
end
# SnakeBridge Ref - normalize to ref wire shape
def encode(%SnakeBridge.Ref{} = ref) do
SnakeBridge.Ref.to_wire_format(ref)
end
# Functions - encode as callback references
def encode(fun) when is_function(fun) do
{:ok, callback_id} = SnakeBridge.CallbackRegistry.register(fun)
arity = Function.info(fun)[:arity]
tagged("callback", %{
"ref_id" => callback_id,
"pid" => self() |> :erlang.pid_to_list() |> IO.iodata_to_binary(),
"arity" => arity
})
end
# Maps - empty map
def encode(%{} = map) when map_size(map) == 0, do: %{}
# Structs that aren't handled above - raise SerializationError
def encode(%{__struct__: _} = struct) do
raise SnakeBridge.SerializationError, value: struct
end
# Maps - check for string keys vs non-string keys
def encode(%{} = map) do
if all_string_keys?(map) do
encode_string_key_map(map)
else
encode_tagged_dict(map)
end
end
# Fallback for unsupported types - raise SerializationError
def encode(other) do
raise SnakeBridge.SerializationError, value: other
end
# Private helpers
defp tagged(type, fields) when is_map(fields) do
fields
|> Map.put("__type__", type)
|> Map.put("__schema__", SnakeBridge.Types.schema_version())
end
defp all_string_keys?(map) do
Enum.all?(map, fn {key, _value} ->
is_binary(key) or (is_atom(key) and key not in [nil, true, false])
end)
end
defp encode_string_key_map(map) do
Map.new(map, fn {key, value} ->
string_key = if is_atom(key), do: Atom.to_string(key), else: key
{string_key, encode(value)}
end)
end
defp encode_tagged_dict(map) do
pairs =
Enum.map(map, fn {key, value} ->
[encode(key), encode(value)]
end)
tagged("dict", %{"pairs" => pairs})
end
defp normalize_pyref_map(ref) do
ref_id = Map.get(ref, :id) || Map.get(ref, :ref_id)
%{}
|> Map.put("__type__", "ref")
|> Map.put("__schema__", SnakeBridge.Ref.schema_version())
|> maybe_put("id", ref_id)
|> maybe_put("session_id", Map.get(ref, :session_id))
|> maybe_put("python_module", Map.get(ref, :python_module))
|> maybe_put("library", Map.get(ref, :library))
end
defp maybe_put(map, _key, nil), do: map
defp maybe_put(map, key, value), do: Map.put(map, key, value)
end
defmodule SnakeBridge.WheelSelector.ConfigStrategy do
@moduledoc false
alias SnakeBridge.WheelConfig
@type wheel_info :: %{
package: String.t(),
version: String.t(),
variant: String.t() | nil,
index_url: String.t() | nil
}
@spec select_wheel(String.t(), String.t(), map()) :: wheel_info()
def select_wheel(package, version, caps) do
variant = variant_for(package, caps)
%{
package: package,
version: version,
variant: variant,
index_url: index_url_for_variant(variant)
}
end
@spec variant_for(String.t(), map()) :: String.t() | nil
def variant_for(package, caps) do
if variant_package?(package) do
variants = WheelConfig.get_variants(package)
cond do
caps.cuda and caps.cuda_version ->
pick_variant(best_cuda_variant(caps.cuda_version), variants)
caps.rocm ->
rocm_variant = WheelConfig.rocm_variant()
pick_variant(rocm_variant, variants)
true ->
pick_variant("cpu", variants)
end
else
nil
end
end
@spec available_variants(String.t()) :: [String.t()]
def available_variants(package) do
if variant_package?(package) do
WheelConfig.get_variants(package)
else
[]
end
end
@spec best_cuda_variant(String.t() | nil) :: String.t()
def best_cuda_variant(nil), do: "cpu"
def best_cuda_variant(cuda_version) do
WheelConfig.get_cuda_mapping(cuda_version) || cuda_variant_fallback(cuda_version)
end
@spec index_url_for_variant(String.t() | nil) :: String.t() | nil
def index_url_for_variant(nil), do: nil
def index_url_for_variant(variant) do
base_url =
Application.get_env(
:snakebridge,
:pytorch_index_base_url,
"https://download.pytorch.org/whl/"
)
"#{String.trim_trailing(base_url, "/")}/#{variant}"
end
defp variant_package?(package) do
package in WheelConfig.packages()
end
defp pick_variant(nil, variants) do
pick_variant("cpu", variants)
end
defp pick_variant(preferred, variants) do
cond do
preferred in variants ->
preferred
"cpu" in variants ->
"cpu"
variants == [] ->
nil
true ->
List.first(variants)
end
end
defp cuda_variant_fallback(version) do
thresholds =
Application.get_env(:snakebridge, :cuda_thresholds, [
{"cu124", 124},
{"cu121", 120},
{"cu118", 117}
])
normalized = normalize_cuda_version(version)
case Integer.parse(normalized || "") do
{value, _} -> find_matching_variant(thresholds, value)
_ -> "cpu"
end
end
defp find_matching_variant(thresholds, cuda_version) do
Enum.find_value(thresholds, "cpu", fn {variant, threshold} ->
if cuda_version >= threshold, do: variant
end)
end
defp normalize_cuda_version(version) when is_binary(version) do
version
|> String.split(".")
|> Enum.take(2)
|> Enum.join()
end
defp normalize_cuda_version(_), do: nil
end
defmodule SnakeBridge.Adapter do
@moduledoc """
Provides the `use SnakeBridge.Adapter` macro for generated Python adapters.
When you `use SnakeBridge.Adapter`, it imports the `__python_call__/2` function
that generated adapters use to call Python functions via Snakepit.
## Example
defmodule MyApp.Math do
use SnakeBridge.Adapter
@spec sqrt(number()) :: float()
def sqrt(x) do
__python_call__("sqrt", [x])
end
end
The adapter module tracks the Python module name and provides the runtime
bridge to execute Python functions.
"""
defmacro __using__(_opts) do
quote do
import SnakeBridge.Adapter, only: [__python_call__: 2]
# Register @python_function as an accumulating attribute for metadata
# This prevents "set but never used" warnings in generated code
Module.register_attribute(__MODULE__, :python_function, accumulate: true)
# Store the Python module name derived from the Elixir module name
@python_module __MODULE__
|> Module.split()
|> List.last()
|> Macro.underscore()
end
end
@doc """
Calls a Python function with the given arguments using SnakeBridge.Runtime.
"""
@spec __python_call__(String.t(), list()) :: {:ok, term()} | {:error, Snakepit.Error.t()}
def __python_call__(func_name, args) do
# Get the calling module to determine the Python module
{module, _func, _arity} =
Process.info(self(), :current_stacktrace)
|> elem(1)
|> Enum.find(fn {mod, _, _, _} ->
mod not in [__MODULE__, Process, :erlang]
end)
|> case do
{mod, func, arity, _} -> {mod, func, arity}
nil -> {nil, nil, nil}
end
if module do
SnakeBridge.Runtime.call(module, func_name, args)
else
{:error, Snakepit.Error.validation_error("Unable to determine calling module", %{})}
end
end
end
defmodule SnakeBridge.Application do
@moduledoc false
use Application
require Logger
@impl true
def start(_type, _args) do
# Suppress noisy logs from OpenTelemetry transitive dependencies
# These apps emit info/notice logs during startup that pollute console output
# Users can override by setting their own log levels in config
suppress_otel_transitive_logs()
children = [
SnakeBridge.SessionManager,
SnakeBridge.CallbackRegistry
]
Supervisor.start_link(children, strategy: :one_for_one, name: SnakeBridge.Supervisor)
end
defp suppress_otel_transitive_logs do
# tls_certificate_check logs "Loading X CA(s) from ..." at notice level
Logger.put_application_level(:tls_certificate_check, :warning)
# opentelemetry_exporter can log during exporter init
Logger.put_application_level(:opentelemetry_exporter, :warning)
end
end
defmodule SnakeBridge.Benchmark do
@moduledoc """
Benchmark utilities for SnakeBridge performance measurement.
Provides functions for measuring execution time, collecting statistics,
and comparing performance across different configurations.
## Usage
# Single measurement
result = Benchmark.measure("my_operation", fn -> do_work() end)
# Multiple iterations with statistics
stats = Benchmark.run_iterations("my_operation", fn -> do_work() end, 10)
# Compare two runs
comparison = Benchmark.compare(baseline_stats, current_stats)
"""
@type measurement :: %{
name: String.t(),
time_us: non_neg_integer(),
value: term(),
error: String.t() | nil
}
@type stats :: %{
name: String.t(),
iterations: non_neg_integer(),
mean_us: float(),
median_us: float(),
min_us: non_neg_integer(),
max_us: non_neg_integer(),
std_dev_us: float(),
times_us: [non_neg_integer()]
}
@type comparison :: %{
speedup: float(),
improvement_percent: float(),
baseline_mean_us: float(),
current_mean_us: float()
}
@doc """
Measures the execution time of a single function call.
Returns a map with:
- `name` - The benchmark name
- `time_us` - Execution time in microseconds
- `value` - The function's return value
- `error` - Error message if the function raised
"""
@spec measure(String.t(), (-> term())) :: measurement()
def measure(name, fun) when is_function(fun, 0) do
{time_us, value} = :timer.tc(fun)
%{
name: name,
time_us: time_us,
value: value,
error: nil
}
rescue
e ->
%{
name: name,
time_us: 0,
value: nil,
error: Exception.message(e)
}
end
@doc """
Runs a function multiple times and collects statistics.
Returns a map with statistical measures:
- `mean_us` - Average time in microseconds
- `median_us` - Median time in microseconds
- `min_us` - Minimum time
- `max_us` - Maximum time
- `std_dev_us` - Standard deviation
"""
@spec run_iterations(String.t(), (-> term()), non_neg_integer()) :: stats()
def run_iterations(name, fun, iterations \\ 10) when is_function(fun, 0) do
# Warmup run
_ = fun.()
times =
Enum.map(1..iterations, fn _ ->
{time_us, _} = :timer.tc(fun)
time_us
end)
mean = Enum.sum(times) / length(times)
sorted = Enum.sort(times)
median = Enum.at(sorted, div(length(sorted), 2))
min = List.first(sorted)
max = List.last(sorted)
std_dev = calculate_std_dev(times, mean)
%{
name: name,
iterations: iterations,
mean_us: mean,
median_us: median,
min_us: min,
max_us: max,
std_dev_us: std_dev,
times_us: times
}
end
@doc """
Compares two benchmark results and calculates improvement metrics.
Returns:
- `speedup` - Ratio (> 1.0 means faster)
- `improvement_percent` - Percentage improvement (positive is faster)
"""
@spec compare(map(), map()) :: comparison()
def compare(%{mean_us: baseline}, %{mean_us: current}) do
speedup = baseline / current
improvement = (1 - current / baseline) * 100
%{
speedup: Float.round(speedup, 3),
improvement_percent: Float.round(improvement, 2),
baseline_mean_us: baseline,
current_mean_us: current
}
end
@doc """
Formats a time in microseconds to a human-readable string.
## Examples
iex> Benchmark.format_time(500)
"500 µs"
iex> Benchmark.format_time(5_000)
"5.00 ms"
iex> Benchmark.format_time(5_000_000)
"5.00 s"
"""
@spec format_time(number()) :: String.t()
def format_time(us) when us < 1_000 do
"#{round(us)} µs"
end
def format_time(us) when us < 1_000_000 do
"#{Float.round(us / 1_000, 2)} ms"
end
def format_time(us) do
"#{Float.round(us / 1_000_000, 2)} s"
end
@doc """
Formats a byte count to a human-readable string.
## Examples
iex> Benchmark.format_bytes(1024)
"1.00 KB"
iex> Benchmark.format_bytes(1_048_576)
"1.00 MB"
"""
@spec format_bytes(number()) :: String.t()
def format_bytes(bytes) when bytes < 1024 do
"#{round(bytes)} B"
end
def format_bytes(bytes) when bytes < 1024 * 1024 do
"#{Float.round(bytes / 1024, 2)} KB"
end
def format_bytes(bytes) when bytes < 1024 * 1024 * 1024 do
"#{Float.round(bytes / (1024 * 1024), 2)} MB"
end
def format_bytes(bytes) do
"#{Float.round(bytes / (1024 * 1024 * 1024), 2)} GB"
end
@doc """
Prints a summary of benchmark statistics.
"""
@spec print_stats(stats()) :: :ok
def print_stats(stats) do
IO.puts("")
IO.puts("#{stats.name}")
IO.puts(String.duplicate("-", String.length(stats.name)))
IO.puts(" Iterations: #{stats.iterations}")
IO.puts(" Mean: #{format_time(stats.mean_us)}")
IO.puts(" Median: #{format_time(stats.median_us)}")
IO.puts(" Min: #{format_time(stats.min_us)}")
IO.puts(" Max: #{format_time(stats.max_us)}")
IO.puts(" Std Dev: #{format_time(stats.std_dev_us)}")
:ok
end
@doc """
Prints a comparison between two benchmark runs.
"""
@spec print_comparison(comparison()) :: :ok
def print_comparison(comparison) do
IO.puts("")
status =
if comparison.improvement_percent >= 0 do
"FASTER"
else
"SLOWER"
end
IO.puts("Comparison: #{status}")
IO.puts(" Speedup: #{comparison.speedup}x")
IO.puts(" Change: #{comparison.improvement_percent}%")
IO.puts(" Baseline: #{format_time(comparison.baseline_mean_us)}")
IO.puts(" Current: #{format_time(comparison.current_mean_us)}")
:ok
end
defp calculate_std_dev(times, mean) do
variance =
times
|> Enum.map(fn t -> :math.pow(t - mean, 2) end)
|> Enum.sum()
|> Kernel./(length(times))
:math.sqrt(variance)
end
end
defmodule SnakeBridge.Bytes do
@moduledoc """
Wrapper struct for binary data that should be sent to Python as `bytes`, not `str`.
By default, SnakeBridge encodes UTF-8 valid Elixir binaries as Python strings.
Use this wrapper when you need to explicitly send data as Python bytes.
## Examples
# Hash a string as bytes
{:ok, hash} = SnakeBridge.call("hashlib", "md5", [SnakeBridge.bytes("abc")])
# Base64 encode
{:ok, encoded} = SnakeBridge.call("base64", "b64encode", [SnakeBridge.bytes("hello")])
# Binary protocol data
{:ok, _} = SnakeBridge.call("struct", "pack", [">I", 42])
## When to Use
Use `SnakeBridge.bytes/1` when calling Python functions that:
- Require `bytes` input (hashlib, cryptography, struct, etc.)
- Work with binary protocols
- Process raw byte data
## Wire Format
Encoded as:
{"__type__": "bytes", "__schema__": 1, "data": ""}
"""
@type t :: %__MODULE__{data: binary()}
defstruct [:data]
@doc """
Creates a Bytes wrapper from binary data.
## Examples
iex> SnakeBridge.Bytes.new("hello")
%SnakeBridge.Bytes{data: "hello"}
iex> SnakeBridge.Bytes.new(<<0, 1, 2, 255>>)
%SnakeBridge.Bytes{data: <<0, 1, 2, 255>>}
"""
@spec new(binary()) :: t()
def new(data) when is_binary(data) do
%__MODULE__{data: data}
end
@doc """
Returns the raw binary data from a Bytes wrapper.
## Examples
iex> bytes = SnakeBridge.Bytes.new("hello")
iex> SnakeBridge.Bytes.data(bytes)
"hello"
"""
@spec data(t()) :: binary()
def data(%__MODULE__{data: data}), do: data
end
defmodule SnakeBridge.CallbackRegistry do
@moduledoc """
Registry for Elixir callbacks passed to Python.
Manages callback lifecycle and provides invocation support.
"""
use GenServer
require Logger
alias SnakeBridge.SessionContext
alias Snakepit.Bridge.ToolRegistry
@tool_name "snakebridge.callback"
@tool_metadata %{
description: "Invoke an Elixir callback from Python",
exposed_to_python: true,
parameters: [
%{name: "callback_id", type: "string", required: true},
%{name: "args", type: "list", required: false}
]
}
# Client API
def start_link(opts \\ []) do
GenServer.start_link(__MODULE__, opts, name: __MODULE__)
end
@doc """
Registers an Elixir function as a callback.
"""
@spec register(function(), pid()) :: {:ok, String.t()}
def register(fun, owner_pid \\ self()) when is_function(fun) do
ensure_tool_registered(current_session_id())
GenServer.call(__MODULE__, {:register, fun, owner_pid})
end
@doc """
Invokes a registered callback with arguments.
"""
@spec invoke(String.t(), list()) :: {:ok, term()} | {:error, term()}
def invoke(callback_id, args) do
GenServer.call(__MODULE__, {:invoke, callback_id, args}, :infinity)
end
@doc """
Unregisters a callback.
"""
@spec unregister(String.t()) :: :ok
def unregister(callback_id) do
GenServer.cast(__MODULE__, {:unregister, callback_id})
end
@doc """
Ensures the callback tool is registered for the session.
"""
@spec ensure_tool_registered(String.t() | nil) :: :ok
def ensure_tool_registered(session_id) do
session_id = session_id || "default"
if Code.ensure_loaded?(ToolRegistry) and
Process.whereis(ToolRegistry) do
case ToolRegistry.register_elixir_tool(
session_id,
@tool_name,
&__MODULE__.handle_tool/1,
@tool_metadata
) do
:ok -> :ok
{:error, {:duplicate_tool, _name}} -> :ok
{:error, _reason} -> :ok
end
end
:ok
end
@doc """
Handles callback tool invocations from Python.
"""
@spec handle_tool(map()) :: map()
def handle_tool(params) when is_map(params) do
callback_id = Map.get(params, "callback_id")
args = params |> Map.get("args", []) |> List.wrap()
decoded_args = Enum.map(args, &SnakeBridge.Types.decode/1)
case invoke(callback_id, decoded_args) do
{:ok, result} ->
SnakeBridge.Types.encode(result)
{:error, reason} ->
%{"__type__" => "callback_error", "reason" => inspect(reason)}
end
end
# Server Implementation
@impl true
def init(_opts) do
state = %{
callbacks: %{},
monitors: %{}
}
{:ok, state}
end
@impl true
def handle_call({:register, fun, owner_pid}, _from, state) do
callback_id = generate_callback_id()
monitor_ref = Process.monitor(owner_pid)
arity = Function.info(fun)[:arity]
callback_data = %{
fun: fun,
owner_pid: owner_pid,
monitor_ref: monitor_ref,
arity: arity
}
new_state = %{
state
| callbacks: Map.put(state.callbacks, callback_id, callback_data),
monitors: Map.put(state.monitors, monitor_ref, callback_id)
}
{:reply, {:ok, callback_id}, new_state}
end
@impl true
def handle_call({:invoke, callback_id, args}, _from, state) do
case Map.get(state.callbacks, callback_id) do
nil ->
{:reply, {:error, :callback_not_found}, state}
%{fun: fun, arity: arity} = _data ->
if length(args) != arity do
{:reply, {:error, {:arity_mismatch, arity}}, state}
else
try do
result = apply(fun, args)
{:reply, {:ok, result}, state}
rescue
exception ->
{:reply, {:error, {:exception, exception}}, state}
end
end
end
end
@impl true
def handle_cast({:unregister, callback_id}, state) do
new_state = do_unregister(state, callback_id)
{:noreply, new_state}
end
@impl true
def handle_info({:DOWN, monitor_ref, :process, _pid, _reason}, state) do
case Map.get(state.monitors, monitor_ref) do
nil ->
{:noreply, state}
callback_id ->
Logger.debug("Callback owner died, unregistering: #{callback_id}")
new_state = do_unregister(state, callback_id)
{:noreply, new_state}
end
end
defp do_unregister(state, callback_id) do
case Map.get(state.callbacks, callback_id) do
nil ->
state
%{monitor_ref: monitor_ref} ->
Process.demonitor(monitor_ref, [:flush])
%{
state
| callbacks: Map.delete(state.callbacks, callback_id),
monitors: Map.delete(state.monitors, monitor_ref)
}
end
end
defp generate_callback_id do
"cb_#{:erlang.unique_integer([:positive])}_#{System.system_time(:millisecond)}"
end
defp current_session_id do
case SessionContext.current() do
%{session_id: session_id} when is_binary(session_id) -> session_id
_ -> "default"
end
end
end
defmodule SnakeBridge.CompileError do
@moduledoc """
Error raised when strict mode detects missing bindings.
"""
defexception [:message]
end
defmodule SnakeBridge.Config do
@moduledoc """
Compile-time configuration for SnakeBridge.
"""
defstruct [
:libraries,
:auto_install,
:generated_dir,
:metadata_dir,
:helper_paths,
:helper_pack_enabled,
:helper_allowlist,
:inline_enabled,
:strict,
:verbose,
:scan_paths,
:scan_exclude,
:introspector,
:docs,
:runtime_client,
:ledger
]
defmodule Library do
@moduledoc """
Configuration struct for a single Python library binding.
"""
defstruct [
:name,
:version,
:module_name,
:python_name,
:pypi_package,
:extras,
include: [],
exclude: [],
streaming: [],
submodules: false
]
@type t :: %__MODULE__{
name: atom(),
version: String.t() | :stdlib | nil,
module_name: module(),
python_name: String.t(),
pypi_package: String.t() | nil,
extras: [String.t()],
include: [String.t()],
exclude: [String.t()],
streaming: [String.t()],
submodules: boolean()
}
end
@type t :: %__MODULE__{
libraries: [Library.t()],
auto_install: :never | :dev | :always,
generated_dir: String.t(),
metadata_dir: String.t(),
helper_paths: [String.t()],
helper_pack_enabled: boolean(),
helper_allowlist: :all | [String.t()],
inline_enabled: boolean(),
strict: boolean(),
verbose: boolean(),
scan_paths: [String.t()],
scan_exclude: [String.t()],
introspector: keyword(),
docs: keyword(),
runtime_client: module(),
ledger: keyword()
}
@doc """
Load config from mix.exs dependency options and Application env.
"""
@spec load() :: t()
def load do
deps = Mix.Project.config()[:deps] || []
opts =
deps
|> Enum.find_value([], fn
{:snakebridge, opts} when is_list(opts) -> opts
{:snakebridge, _req, opts} when is_list(opts) -> opts
_ -> nil
end)
|> List.wrap()
%__MODULE__{
libraries: parse_libraries(Keyword.get(opts, :libraries, [])),
auto_install: Application.get_env(:snakebridge, :auto_install, :dev),
generated_dir: Keyword.get(opts, :generated_dir, "lib/snakebridge_generated"),
metadata_dir: Keyword.get(opts, :metadata_dir, ".snakebridge"),
helper_paths: Application.get_env(:snakebridge, :helper_paths, ["priv/python/helpers"]),
helper_pack_enabled: Application.get_env(:snakebridge, :helper_pack_enabled, true),
helper_allowlist: Application.get_env(:snakebridge, :helper_allowlist, :all),
inline_enabled: Application.get_env(:snakebridge, :inline_enabled, false),
strict: env_flag(:strict, "SNAKEBRIDGE_STRICT", false),
verbose: env_flag(:verbose, "SNAKEBRIDGE_VERBOSE", false),
scan_paths: Application.get_env(:snakebridge, :scan_paths, ["lib"]),
scan_exclude: Application.get_env(:snakebridge, :scan_exclude, []),
introspector: Application.get_env(:snakebridge, :introspector, []),
docs: Application.get_env(:snakebridge, :docs, []),
runtime_client: Application.get_env(:snakebridge, :runtime_client, Snakepit),
ledger: Application.get_env(:snakebridge, :ledger, [])
}
end
@doc false
def parse_libraries(libraries) when is_list(libraries) do
Enum.map(libraries, &parse_library/1)
end
defp parse_library({name, version}) when is_binary(version) or version == :stdlib do
build_library(name, version, [])
end
defp parse_library({name, opts}) when is_list(opts) do
version = Keyword.get(opts, :version)
build_library(name, version, opts)
end
defp parse_library(name) when is_atom(name) do
build_library(name, nil, [])
end
defp parse_library(name) when is_binary(name) do
build_library(String.to_atom(name), nil, [])
end
defp build_library(name, version, opts) do
module_name = Keyword.get(opts, :module_name, default_module_name(name))
python_name = Keyword.get(opts, :python_name, Atom.to_string(name))
extras = Keyword.get(opts, :extras, [])
%Library{
name: name,
version: version,
module_name: module_name,
python_name: python_name,
pypi_package: Keyword.get(opts, :pypi_package),
extras: List.wrap(extras),
include: Keyword.get(opts, :include, []),
exclude: Keyword.get(opts, :exclude, []),
streaming: Keyword.get(opts, :streaming, []),
submodules: Keyword.get(opts, :submodules, false)
}
end
defp default_module_name(name) do
name
|> Atom.to_string()
|> Macro.camelize()
|> then(&Module.concat([&1]))
end
defp env_flag(config_key, env_var, default) do
case System.get_env(env_var) do
nil -> Application.get_env(:snakebridge, config_key, default)
value -> value in ["1", "true", "TRUE", "yes", "YES"]
end
end
end
defmodule SnakeBridge.Defaults do
@moduledoc """
Centralized defaults for all configurable values in SnakeBridge.
All values can be overridden via `Application.get_env(:snakebridge, key)`.
## Configuration Options
### Introspection
- `:introspector_timeout` - Timeout in ms for introspecting Python modules (default: `30_000`)
- `:introspector_max_concurrency` - Max concurrent introspection tasks (default: `System.schedulers_online()`)
### Wheel Selector (PyTorch/CUDA)
- `:pytorch_index_base_url` - Base URL for PyTorch wheel index (default: `"https://download.pytorch.org/whl/"`)
- `:cuda_thresholds` - CUDA version to variant mapping (default: `[{"cu124", 124}, {"cu121", 120}, {"cu118", 117}]`)
### Session Lifecycle
- `:session_max_refs` - Maximum refs per session (default: `10_000`)
- `:session_ttl_seconds` - Session time-to-live in seconds (default: `3600`)
### Code Generation
- `:variadic_max_arity` - Max arity for variadic wrappers (default: `8`)
- `:generated_dir` - Directory for generated code (default: `"lib/snakebridge_generated"`)
- `:metadata_dir` - Directory for metadata files (default: `".snakebridge"`)
### Protocol
- `:protocol_version` - Wire protocol version (default: `1`)
- `:min_supported_version` - Minimum supported protocol version (default: `1`)
### Runtime Timeouts
Runtime timeout configuration is nested under the `:runtime` key:
- `:timeout_profile` - Default profile for calls (default: `:default` for calls, `:streaming` for streams)
- `:default_timeout` - Default unary call timeout in ms (default: `120_000`)
- `:default_stream_timeout` - Default stream timeout in ms (default: `1_800_000`)
- `:library_profiles` - Map of library names to profiles (default: `%{}`)
- `:profiles` - Map of profile names to timeout settings
Built-in profiles:
- `:default` - 120s timeout for regular calls
- `:streaming` - 120s timeout, 30min stream_timeout
- `:ml_inference` - 10min timeout for ML/LLM workloads
- `:batch_job` - infinity timeout for long-running jobs
## Example Configuration
config :snakebridge,
introspector_timeout: 60_000,
pytorch_index_base_url: "https://my-mirror.example.com/pytorch/",
cuda_thresholds: [
{"cu126", 126},
{"cu124", 124},
{"cu121", 120},
{"cu118", 117}
],
session_max_refs: 50_000,
session_ttl_seconds: 7200,
runtime: [
timeout_profile: :default,
library_profiles: %{
"transformers" => :ml_inference,
"torch" => :batch_job
},
profiles: %{
default: [timeout: 120_000],
ml_inference: [timeout: 600_000, stream_timeout: 1_800_000],
batch_job: [timeout: :infinity, stream_timeout: :infinity]
}
]
"""
# Introspection
def introspector_timeout, do: get(:introspector_timeout, 30_000)
def introspector_max_concurrency,
do: get(:introspector_max_concurrency, System.schedulers_online())
# Wheel selector
def pytorch_index_base_url,
do: get(:pytorch_index_base_url, "https://download.pytorch.org/whl/")
def cuda_thresholds do
get(:cuda_thresholds, [
{"cu124", 124},
{"cu121", 120},
{"cu118", 117}
])
end
# Session context
def session_max_refs, do: get(:session_max_refs, 10_000)
def session_ttl_seconds, do: get(:session_ttl_seconds, 3600)
# Protocol
def protocol_version, do: get(:protocol_version, 1)
def min_supported_version, do: get(:min_supported_version, 1)
# Code generation
def variadic_max_arity, do: get(:variadic_max_arity, 8)
def generated_dir, do: get(:generated_dir, "lib/snakebridge_generated")
def metadata_dir, do: get(:metadata_dir, ".snakebridge")
# ============================================================================
# Runtime Timeout Configuration
# ============================================================================
@default_runtime_profiles %{
default: [timeout: 120_000],
streaming: [timeout: 120_000, stream_timeout: 1_800_000],
ml_inference: [timeout: 600_000, stream_timeout: 1_800_000],
batch_job: [timeout: :infinity, stream_timeout: :infinity]
}
@doc """
Returns the runtime configuration keyword list.
"""
@spec runtime_config() :: keyword()
def runtime_config, do: Application.get_env(:snakebridge, :runtime, [])
@doc """
Returns the timeout profile for a given call kind.
Call kinds:
- `:call` - Regular function calls (default: `:default`)
- `:stream` - Streaming calls (default: `:streaming`)
"""
@spec runtime_timeout_profile(atom()) :: atom()
def runtime_timeout_profile(call_kind \\ :call) do
runtime_config()
|> Keyword.get(:timeout_profile, default_timeout_profile(call_kind))
end
@doc """
Returns configured library-to-profile mappings.
Example:
config :snakebridge, runtime: [
library_profiles: %{
"transformers" => :ml_inference,
"torch" => :batch_job
}
]
"""
@spec runtime_library_profiles() :: map()
def runtime_library_profiles do
runtime_config() |> Keyword.get(:library_profiles, %{})
end
@doc """
Returns all timeout profiles.
Default profiles:
- `:default` - 120s timeout for regular calls
- `:streaming` - 120s timeout, 30min stream_timeout
- `:ml_inference` - 10min timeout for ML/LLM workloads
- `:batch_job` - infinity timeout for long-running jobs
"""
@spec runtime_profiles() :: map()
def runtime_profiles do
runtime_config() |> Keyword.get(:profiles, @default_runtime_profiles)
end
@doc """
Returns the default unary call timeout in milliseconds.
"""
@spec runtime_default_timeout() :: timeout()
def runtime_default_timeout do
runtime_config() |> Keyword.get(:default_timeout, 120_000)
end
@doc """
Returns the default stream timeout in milliseconds.
"""
@spec runtime_default_stream_timeout() :: timeout()
def runtime_default_stream_timeout do
runtime_config() |> Keyword.get(:default_stream_timeout, 1_800_000)
end
defp default_timeout_profile(:stream), do: :streaming
defp default_timeout_profile(_), do: :default
@doc """
Returns all current configuration values as a map.
"""
@spec all() :: map()
def all do
%{
introspector_timeout: introspector_timeout(),
introspector_max_concurrency: introspector_max_concurrency(),
pytorch_index_base_url: pytorch_index_base_url(),
cuda_thresholds: cuda_thresholds(),
session_max_refs: session_max_refs(),
session_ttl_seconds: session_ttl_seconds(),
protocol_version: protocol_version(),
min_supported_version: min_supported_version(),
variadic_max_arity: variadic_max_arity(),
generated_dir: generated_dir(),
metadata_dir: metadata_dir(),
runtime_default_timeout: runtime_default_timeout(),
runtime_default_stream_timeout: runtime_default_stream_timeout(),
runtime_timeout_profile: runtime_timeout_profile()
}
end
defp get(key, default) do
Application.get_env(:snakebridge, key, default)
end
end
defmodule SnakeBridge.Docs do
@moduledoc """
On-demand documentation fetching with optional caching.
"""
@cache_table :snakebridge_docs
@spec get(module(), atom() | String.t()) :: String.t()
def get(module, function) do
start_time = System.monotonic_time()
key = {module, function}
case lookup_cache(key) do
{:hit, doc} ->
SnakeBridge.Telemetry.docs_fetch(start_time, module, function, :cache)
doc
:miss ->
{doc, source} = fetch_doc_with_source(module, function)
maybe_cache(key, doc)
SnakeBridge.Telemetry.docs_fetch(start_time, module, function, source)
doc
end
end
defp fetch_doc_with_source(module, function) do
case docs_source() do
:python ->
{fetch_from_python(module, function), :python}
:metadata ->
{fetch_from_metadata(module, function) || "Documentation unavailable.", :metadata}
:hybrid ->
fetch_hybrid_doc(module, function)
end
end
defp fetch_hybrid_doc(module, function) do
case fetch_from_metadata(module, function) do
nil -> {fetch_from_python(module, function), :python}
metadata -> {metadata, :metadata}
end
end
@spec search(module(), String.t()) :: list()
def search(module, query) when is_binary(query) do
query = query |> String.trim() |> String.downcase()
module
|> functions_for_search()
|> Enum.map(fn {name, summary} -> {name, summary, score(name, query)} end)
|> Enum.filter(fn {_name, _summary, relevance} -> relevance > 0.3 end)
|> Enum.sort_by(fn {_name, _summary, relevance} -> -relevance end)
|> Enum.take(10)
|> Enum.map(fn {name, summary, relevance} ->
%{name: name, summary: summary, relevance: relevance}
end)
end
defp docs_source do
Application.get_env(:snakebridge, :docs, [])
|> Keyword.get(:source, :python)
end
defp functions_for_search(module) do
if function_exported?(module, :__functions__, 0) do
module.__functions__()
|> Enum.map(fn {name, _arity, _mod, summary} ->
{name, summary |> to_string()}
end)
else
[]
end
end
defp score(name, query) do
name = name |> to_string() |> String.downcase()
cond do
query == "" -> 0.0
name == query -> 1.0
String.starts_with?(name, query) -> 0.9
String.contains?(name, query) -> 0.7
true -> 0.0
end
end
defp fetch_from_metadata(module, function) do
function_name = to_string(function)
with {:docs_v1, _, _, _, _, _, docs} <- Code.fetch_docs(module),
entry when not is_nil(entry) <- Enum.find(docs, &doc_entry_matches?(&1, function_name)),
docstring when is_binary(docstring) <- docstring_from_entry(entry) do
normalize_docstring(docstring)
else
_ -> nil
end
end
defp doc_entry_matches?({{kind, name, _arity}, _, _, _, _}, function_name)
when kind in [:function, :macro] and is_atom(name) do
Atom.to_string(name) == function_name
end
defp doc_entry_matches?({{name, _arity}, _, _, _, _}, function_name) when is_atom(name) do
Atom.to_string(name) == function_name
end
defp doc_entry_matches?(_entry, _function_name), do: false
defp docstring_from_entry({_id, _anno, _signature, doc, _metadata}) do
case doc do
:hidden -> nil
:none -> nil
{_, text} when is_binary(text) -> text
text when is_binary(text) -> text
_ -> nil
end
end
defp normalize_docstring(docstring) do
case String.trim(docstring) do
"" -> nil
trimmed -> trimmed
end
end
defp fetch_from_python(module, function) do
python_name = python_module_name(module)
script = doc_script()
case python_runner().run(script, [python_name, to_string(function)], []) do
{:ok, output} -> String.trim(output)
{:error, _} -> "Documentation unavailable."
end
end
defp python_runner do
Application.get_env(:snakebridge, :python_runner, SnakeBridge.PythonRunner.System)
end
defp python_module_name(module) do
if function_exported?(module, :__snakebridge_python_name__, 0) do
module.__snakebridge_python_name__()
else
module
|> Module.split()
|> Enum.map_join(".", &Macro.underscore/1)
end
end
defp doc_script do
~S"""
import importlib
import inspect
import sys
module_name = sys.argv[1]
function_name = sys.argv[2]
module = importlib.import_module(module_name)
obj = getattr(module, function_name, None)
if obj is None:
print("Function not found.")
else:
print(inspect.getdoc(obj) or "Documentation unavailable.")
"""
end
defp lookup_cache(key) do
if cache_enabled?() do
ensure_cache_table()
case :ets.lookup(@cache_table, key) do
[{^key, doc}] -> {:hit, doc}
[] -> :miss
end
else
:miss
end
end
defp maybe_cache(key, doc) do
if cache_enabled?() do
ensure_cache_table()
:ets.insert(@cache_table, {key, doc})
end
end
defp cache_enabled? do
Application.get_env(:snakebridge, :docs, [])
|> Keyword.get(:cache_enabled, true)
end
defp ensure_cache_table do
case :ets.whereis(@cache_table) do
:undefined ->
:ets.new(@cache_table, [:set, :public, :named_table, read_concurrency: true])
_ ->
@cache_table
end
end
end
defmodule SnakeBridge.DynamicException do
@moduledoc """
Dynamically creates Elixir exception modules from Python exception class names.
This enables pattern matching on Python exceptions:
rescue
e in SnakeBridge.DynamicException.ValueError ->
handle_value_error(e)
"""
@exception_cache :snakebridge_exception_cache
@doc """
Creates an exception struct from a Python class name and message.
"""
@spec create(String.t(), String.t() | nil, keyword()) :: Exception.t()
def create(python_class_name, message, opts \\ []) when is_binary(python_class_name) do
module = get_or_create_module(python_class_name)
details =
opts
|> Keyword.delete(:python_traceback)
|> then(fn cleaned -> Keyword.get(cleaned, :details, cleaned) end)
struct(module,
message: message || "",
python_class: python_class_name,
details: details,
python_traceback: Keyword.get(opts, :python_traceback)
)
end
@doc """
Gets or creates an exception module for a Python class name.
"""
@spec get_or_create_module(String.t()) :: module()
def get_or_create_module(python_class_name) when is_binary(python_class_name) do
ensure_cache_exists()
class_name = sanitize_class_name(python_class_name)
module_name = Module.concat(__MODULE__, class_name)
if Code.ensure_loaded?(module_name) do
:ets.insert(@exception_cache, {module_name, true})
module_name
else
case :ets.lookup(@exception_cache, module_name) do
[{^module_name, true}] ->
module_name
[] ->
create_exception_module(module_name, python_class_name)
module_name
end
end
end
@doc false
def ensure_cache_exists do
if :ets.whereis(@exception_cache) == :undefined do
:ets.new(@exception_cache, [:named_table, :set, :public])
end
end
defp create_exception_module(module_name, python_class_name) do
unless Code.ensure_loaded?(module_name) do
Module.create(
module_name,
quote do
@moduledoc """
Dynamic exception for Python `#{unquote(python_class_name)}`.
"""
defexception [:message, :python_class, :details, :python_traceback]
@impl true
def message(%{message: message}), do: message || ""
end,
Macro.Env.location(__ENV__)
)
end
:ets.insert(@exception_cache, {module_name, true})
end
defp sanitize_class_name(python_class_name) do
python_class_name
|> String.split(".")
|> List.last()
|> String.replace(~r/[^A-Za-z0-9_]/, "")
|> String.split("_", trim: true)
|> Enum.map_join("", &capitalize_preserve/1)
|> normalize_name()
end
defp normalize_name(""), do: "PythonError"
defp normalize_name(<> = name) when first in ?0..?9 do
"Py" <> name
end
defp normalize_name(name), do: name
defp capitalize_preserve(<>) when first in ?a..?z do
<> <> rest
end
defp capitalize_preserve(segment), do: segment
end
defmodule SnakeBridge.Dynamic do
@moduledoc """
Dynamic dispatch for calling methods on Python objects without generated code.
Use this module when:
- Python returns an object of a class you did not generate bindings for
- You need to call methods dynamically at runtime
- You want a no-codegen escape hatch for refs
"""
alias SnakeBridge.Runtime
@type ref :: SnakeBridge.Ref.t() | map()
@type opts :: keyword()
@doc """
Calls a method on a Python object reference.
"""
@spec call(ref(), atom() | String.t(), list(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def call(ref, method, args \\ [], opts \\ []) do
validate_ref!(ref)
Runtime.call_method(ref, method, args, opts)
end
@doc """
Gets an attribute from a Python object reference.
"""
@spec get_attr(ref(), atom() | String.t(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def get_attr(ref, attr, opts \\ []) do
validate_ref!(ref)
Runtime.get_attr(ref, attr, opts)
end
@doc """
Sets an attribute on a Python object reference.
"""
@spec set_attr(ref(), atom() | String.t(), term(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def set_attr(ref, attr, value, opts \\ []) do
validate_ref!(ref)
Runtime.set_attr(ref, attr, value, opts)
end
@doc """
Checks if a value is a valid Python reference.
"""
@spec ref?(term()) :: boolean()
def ref?(value), do: SnakeBridge.Ref.ref?(value)
@doc false
def build_call_payload(ref, method, args) do
wire_ref = SnakeBridge.Ref.to_wire_format(ref)
%{
"call_type" => "method",
"instance" => wire_ref,
"method" => to_string(method),
"args" => args
}
end
defp validate_ref!(ref) do
unless ref?(ref) do
raise ArgumentError, "Invalid ref: expected a SnakeBridge ref, got: #{inspect(ref)}"
end
end
end
defmodule SnakeBridge.EnvironmentError do
@moduledoc """
Error raised when required Python packages are missing.
"""
defexception [:message, :missing_packages, :suggestion]
@type t :: %__MODULE__{
message: String.t(),
missing_packages: [String.t()],
suggestion: String.t()
}
@impl Exception
def message(%__MODULE__{message: message, suggestion: suggestion}) do
if suggestion do
message <> "\n\nSuggestion: " <> suggestion
else
message
end
end
end
defmodule SnakeBridge.ErrorTranslator do
@moduledoc """
Translates Python/ML errors into structured SnakeBridge errors.
This module recognizes common ML error patterns from PyTorch, NumPy,
and other ML libraries, and translates them into structured error
types with actionable suggestions.
## Supported Error Types
- `SnakeBridge.Error.ShapeMismatchError` - Tensor shape incompatibilities
- `SnakeBridge.Error.OutOfMemoryError` - GPU/CPU memory exhaustion
- `SnakeBridge.Error.DtypeMismatchError` - Tensor dtype incompatibilities
## Examples
iex> error = %RuntimeError{message: "CUDA out of memory"}
iex> SnakeBridge.ErrorTranslator.translate(error)
%SnakeBridge.Error.OutOfMemoryError{device: {:cuda, 0}, ...}
"""
alias SnakeBridge.DynamicException
alias SnakeBridge.Error.{DtypeMismatchError, OutOfMemoryError, ShapeMismatchError}
alias SnakeBridge.{InvalidRefError, RefNotFoundError, SessionMismatchError}
# Mapping from normalized dtype strings to Elixir atoms
@dtype_map %{
# PyTorch short names
"float" => :float32,
"double" => :float64,
"half" => :float16,
"long" => :int64,
"int" => :int32,
"short" => :int16,
"byte" => :uint8,
"char" => :int8,
"bool" => :bool,
# PyTorch qualified names
"torch.float32" => :float32,
"torch.float64" => :float64,
"torch.float16" => :float16,
"torch.int64" => :int64,
"torch.int32" => :int32,
"torch.int16" => :int16,
"torch.int8" => :int8,
"torch.uint8" => :uint8,
"torch.bool" => :bool,
"torch.bfloat16" => :bfloat16
}
@doc """
Translates a Python/ML error to a structured SnakeBridge error.
Returns the original error if it cannot be translated.
"""
@spec translate(Exception.t() | map() | nil, String.t() | nil) :: Exception.t() | nil
def translate(error, traceback \\ nil)
def translate(nil, _traceback), do: nil
def translate(%RuntimeError{message: message} = error, traceback) do
case translate_message(message) do
nil -> error
translated -> maybe_add_traceback(translated, traceback)
end
end
def translate(%{python_type: type} = error, traceback) when is_binary(type) do
translate_python_error(type, error, traceback)
end
def translate(%{"python_type" => type} = error, traceback) when is_binary(type) do
translate_python_error(type, error, traceback)
end
def translate(%{error_type: type} = error, traceback) when is_binary(type) do
translate_python_error(type, error, traceback)
end
def translate(%{"error_type" => type} = error, traceback) when is_binary(type) do
translate_python_error(type, error, traceback)
end
def translate(error, _traceback), do: error
@doc """
Translates an error message string to a structured error.
Returns nil if the message cannot be translated.
"""
@spec translate_message(String.t()) :: Exception.t() | nil
def translate_message(message) when is_binary(message) do
cond do
ref_not_found?(message) -> translate_ref_not_found(message)
session_mismatch?(message) -> translate_session_mismatch(message)
invalid_ref?(message) -> translate_invalid_ref(message)
shape_mismatch?(message) -> translate_shape_error(message)
oom_error?(message) -> translate_oom_error(message)
dtype_mismatch?(message) -> translate_dtype_error(message)
true -> nil
end
end
defp translate_python_error(type, error, traceback) do
message = error_message(error)
case maybe_translate_message(message, traceback) do
{:ok, translated} ->
translated
:error ->
DynamicException.create(type, message, details: error, python_traceback: traceback)
end
end
defp maybe_translate_message(message, traceback) when is_binary(message) do
case translate_message(message) do
nil -> :error
translated -> {:ok, maybe_add_traceback(translated, traceback)}
end
end
defp maybe_translate_message(_message, _traceback), do: :error
defp error_message(error) when is_map(error) do
value =
Map.get(error, :message) ||
Map.get(error, "message") ||
Map.get(error, :error) ||
Map.get(error, "error")
cond do
is_binary(value) -> value
is_nil(value) -> ""
true -> inspect(value)
end
end
@doc """
Converts a Python/PyTorch dtype string to an Elixir atom.
## Examples
iex> SnakeBridge.ErrorTranslator.dtype_from_string("Float")
:float32
iex> SnakeBridge.ErrorTranslator.dtype_from_string("torch.float64")
:float64
"""
@spec dtype_from_string(String.t()) :: atom()
def dtype_from_string(dtype_str) do
normalized = dtype_str |> String.trim() |> String.downcase()
case Map.fetch(@dtype_map, normalized) do
{:ok, dtype} -> dtype
:error -> String.to_atom(String.replace(normalized, ".", "_"))
end
end
# Shape mismatch detection patterns
defp shape_mismatch?(message) do
String.contains?(message, "shapes cannot be multiplied") or
String.contains?(message, "size of tensor") or
String.contains?(message, "incompatible shapes") or
String.contains?(message, "Dimension out of range") or
String.contains?(message, "shape mismatch") or
String.contains?(message, "dimension mismatch")
end
# OOM error detection patterns
defp oom_error?(message) do
String.contains?(message, "out of memory") or
String.contains?(message, "OutOfMemory") or
String.contains?(message, "OOM")
end
# Dtype mismatch detection patterns
defp dtype_mismatch?(message) do
String.contains?(message, "expected scalar type") or
String.contains?(message, "expected dtype") or
String.contains?(message, "type mismatch")
end
# Ref lifecycle error detection patterns
defp ref_not_found?(message) do
String.contains?(message, "Unknown SnakeBridge reference")
end
defp session_mismatch?(message) do
String.contains?(message, "SnakeBridge reference session mismatch")
end
defp invalid_ref?(message) do
String.contains?(message, "Invalid SnakeBridge reference") or
String.contains?(message, "SnakeBridge reference missing id")
end
# Translate shape errors
defp translate_shape_error(message) do
cond do
# mat1 and mat2 shapes cannot be multiplied (3x4 and 5x6)
match = Regex.run(~r/shapes cannot be multiplied \((\d+)x(\d+) and (\d+)x(\d+)\)/, message) ->
[_, a_rows, a_cols, b_rows, b_cols] = match
ShapeMismatchError.new(:matmul,
shape_a: [String.to_integer(a_rows), String.to_integer(a_cols)],
shape_b: [String.to_integer(b_rows), String.to_integer(b_cols)],
message: extract_core_message(message)
)
# Broadcasting shape mismatch
String.contains?(message, "broadcasting") ->
shapes = extract_broadcast_shapes(message)
ShapeMismatchError.new(:broadcast,
shape_a: elem(shapes, 0),
shape_b: elem(shapes, 1),
message: extract_core_message(message)
)
# Dimension errors
String.contains?(message, "Dimension") ->
ShapeMismatchError.new(:index,
message: extract_core_message(message)
)
# Size mismatch
String.contains?(message, "size of tensor") ->
ShapeMismatchError.new(:elementwise,
message: extract_core_message(message)
)
# Generic shape error
true ->
ShapeMismatchError.new(:unknown,
message: extract_core_message(message)
)
end
end
# Translate OOM errors
defp translate_oom_error(message) do
device = detect_device(message)
memory_info = extract_memory_info(message)
OutOfMemoryError.new(device,
requested_mb: memory_info[:requested],
available_mb: memory_info[:available],
total_mb: memory_info[:total],
message: extract_core_message(message)
)
end
# Translate dtype errors
defp translate_dtype_error(message) do
{expected, got} = extract_dtype_info(message)
DtypeMismatchError.new(expected, got, message: extract_core_message(message))
end
# Translate ref not found errors
defp translate_ref_not_found(message) do
ref_id = extract_ref_id(message)
RefNotFoundError.exception(
ref_id: ref_id,
message: message
)
end
# Translate session mismatch errors
defp translate_session_mismatch(message) do
SessionMismatchError.exception(message: message)
end
# Translate invalid ref errors
defp translate_invalid_ref(message) do
reason = extract_invalid_reason(message)
InvalidRefError.exception(
reason: reason,
message: message
)
end
# Extract ref ID from error message
defp extract_ref_id(message) do
case Regex.run(~r/reference[:\s]+['\"]?([a-zA-Z0-9_]+)['\"]?/i, message) do
[_, ref_id] -> ref_id
nil -> nil
end
end
# Extract invalid ref reason
defp extract_invalid_reason(message) do
cond do
String.contains?(message, "missing id") -> :missing_id
String.contains?(message, "payload") -> :invalid_format
true -> :unknown
end
end
# Device detection from error message
defp detect_device(message) do
cond do
String.contains?(message, "MPS") -> :mps
match = Regex.run(~r/GPU (\d+)/, message) -> {:cuda, String.to_integer(Enum.at(match, 1))}
String.contains?(message, "CUDA") -> {:cuda, 0}
String.contains?(message, "cuda") -> {:cuda, 0}
true -> :cpu
end
end
# Extract memory information from OOM message
defp extract_memory_info(message) do
requested = extract_memory_value(message, ~r/allocate (\d+) MiB/)
total = extract_memory_value(message, ~r/(\d+) MiB total/)
available = extract_memory_value(message, ~r/(\d+) MiB free/)
%{requested: requested, total: total, available: available}
end
defp extract_memory_value(message, pattern) do
case Regex.run(pattern, message) do
[_, value] -> String.to_integer(value)
nil -> nil
end
end
# Extract dtype info from error message
defp extract_dtype_info(message) do
cond do
# "expected scalar type Float but found Double"
match = Regex.run(~r/expected scalar type (\w+) but found (\w+)/, message) ->
[_, expected, got] = match
{dtype_from_string(expected), dtype_from_string(got)}
# "expected dtype torch.float32 but got torch.int64"
match = Regex.run(~r/expected dtype ([\w.]+) but got ([\w.]+)/, message) ->
[_, expected, got] = match
{dtype_from_string(expected), dtype_from_string(got)}
true ->
{:unknown, :unknown}
end
end
# Extract broadcast shapes from message
defp extract_broadcast_shapes(message) do
case Regex.run(~r/\[([^\]]+)\] vs \[([^\]]+)\]/, message) do
[_, a, b] ->
{parse_shape(a), parse_shape(b)}
nil ->
{nil, nil}
end
end
defp parse_shape(shape_str) do
shape_str
|> String.split(",")
|> Enum.map(&String.trim/1)
|> Enum.map(&String.to_integer/1)
end
# Extract the core error message
defp extract_core_message(message) do
message
|> String.trim()
|> String.replace(~r/^RuntimeError:\s*/, "")
|> String.replace(~r/^torch\.\w+Error:\s*/, "")
|> String.split("\n")
|> List.first()
|> String.trim()
end
# Add traceback if provided
defp maybe_add_traceback(error, nil), do: error
defp maybe_add_traceback(%{__struct__: _} = error, traceback) do
Map.put(error, :python_traceback, traceback)
end
end
defmodule SnakeBridge.Error do
@moduledoc """
ML-specific error types for SnakeBridge.
This module provides structured error types that translate Python/ML
errors into Elixir exceptions with actionable suggestions.
## Available Error Types
- `SnakeBridge.Error.ShapeMismatchError` - Tensor shape incompatibilities
- `SnakeBridge.Error.OutOfMemoryError` - GPU/CPU memory exhaustion
- `SnakeBridge.Error.DtypeMismatchError` - Tensor dtype incompatibilities
## Translation
Use `SnakeBridge.ErrorTranslator` to automatically translate Python
exceptions into these structured error types.
## Examples
# Creating errors directly
error = SnakeBridge.Error.ShapeMismatchError.new(:matmul,
shape_a: [3, 4],
shape_b: [5, 6]
)
# Raising errors
raise SnakeBridge.Error.OutOfMemoryError, device: {:cuda, 0}
# Translating Python errors
translated = SnakeBridge.ErrorTranslator.translate(python_error)
"""
# Re-export error modules for convenient access
defdelegate shape_mismatch(operation, opts \\ []), to: __MODULE__.ShapeMismatchError, as: :new
defdelegate out_of_memory(device, opts \\ []), to: __MODULE__.OutOfMemoryError, as: :new
defdelegate dtype_mismatch(expected, got, opts \\ []),
to: __MODULE__.DtypeMismatchError,
as: :new
end
defmodule SnakeBridge.Examples do
@moduledoc false
@failure_key :snakebridge_example_failures
def reset_failures do
Process.put(@failure_key, 0)
end
def record_failure do
Process.put(@failure_key, failure_count() + 1)
end
def failure_count do
Process.get(@failure_key, 0)
end
def assert_no_failures! do
count = failure_count()
if count > 0 do
raise "Example failed with #{count} unexpected error(s)."
end
end
def assert_script_ok(result) do
case result do
{:error, reason} ->
raise "Snakepit script failed: #{inspect(reason)}"
_ ->
:ok
end
end
end
defmodule SnakeBridge.Generator do
@moduledoc """
Generates Elixir source files from introspection data.
"""
alias SnakeBridge.Docs.{MarkdownConverter, RstParser}
alias SnakeBridge.Generator.TypeMapper
@reserved_words ~w(def defp defmodule class do end if unless case cond for while with fn when and or not true false nil in try catch rescue after else raise throw receive)
@dunder_mappings %{
"__init__" => "new",
"__str__" => "to_string",
"__repr__" => "inspect",
"__len__" => "length",
"__getitem__" => "get",
"__setitem__" => "put",
"__contains__" => "member?"
}
@spec render_library(SnakeBridge.Config.Library.t(), list(), list(), keyword()) :: String.t()
def render_library(library, functions, classes, opts \\ []) do
version = Keyword.get(opts, :version, Application.spec(:snakebridge, :vsn) |> to_string())
module_name = module_to_string(library.module_name)
header = """
# Generated by SnakeBridge v#{version} - DO NOT EDIT MANUALLY
# Regenerate with: mix compile
# Library: #{library.python_name} #{library.version || "unknown"}
"""
moduledoc = """
@moduledoc \"\"\"
SnakeBridge bindings for `#{library.python_name}`.
## Runtime Options
All functions accept a `__runtime__` option for controlling execution behavior:
#{library.module_name}.some_function(args, __runtime__: [timeout: 120_000])
### Supported runtime options
- `:timeout` - Call timeout in milliseconds (default: 120,000ms / 2 minutes)
- `:timeout_profile` - Use a named profile (`:default`, `:ml_inference`, `:batch_job`, `:streaming`)
- `:stream_timeout` - Timeout for streaming operations (default: 1,800,000ms / 30 minutes)
- `:session_id` - Override the session ID for this call
### Timeout Profiles
- `:default` - 2 minute timeout for regular calls
- `:ml_inference` - 10 minute timeout for ML/LLM workloads
- `:batch_job` - Unlimited timeout for long-running jobs
- `:streaming` - 2 minute timeout, 30 minute stream_timeout
### Example with timeout override
# For a long-running ML inference call
#{library.module_name}.predict(data, __runtime__: [timeout_profile: :ml_inference])
# Or explicit timeout
#{library.module_name}.predict(data, __runtime__: [timeout: 600_000])
See `SnakeBridge.Defaults` for global timeout configuration.
\"\"\"
def __snakebridge_python_name__, do: "#{library.python_name}"
def __snakebridge_library__, do: "#{library.python_name}"
"""
functions_by_module =
functions
|> Enum.map(&Map.put_new(&1, "python_module", library.python_name))
|> Enum.group_by(& &1["python_module"])
base_functions = Map.get(functions_by_module, library.python_name, [])
function_defs =
base_functions
|> Enum.sort_by(& &1["name"])
|> Enum.map_join("\n\n", &render_function(&1, library))
submodule_defs =
functions_by_module
|> Map.drop([library.python_name])
|> Enum.sort_by(fn {python_module, _} -> python_module end)
|> Enum.map_join("\n\n", fn {python_module, funcs} ->
render_submodule(python_module, funcs, library)
end)
class_defs =
classes
|> Enum.sort_by(&class_sort_key/1)
|> Enum.map_join("\n\n", &render_class(&1, library))
discovery = render_discovery(functions, classes)
"""
#{header}
defmodule #{module_name} do
#{moduledoc}
#{function_defs}
#{submodule_defs}
#{class_defs}
#{discovery}
end
"""
|> Code.format_string!()
|> IO.iodata_to_binary()
end
@spec generate_library(SnakeBridge.Config.Library.t(), list(), list(), SnakeBridge.Config.t()) ::
:ok
def generate_library(library, functions, classes, config) do
start_time = System.monotonic_time()
File.mkdir_p!(config.generated_dir)
path = Path.join(config.generated_dir, "#{library.python_name}.ex")
source =
render_library(library, functions, classes, version: Application.spec(:snakebridge, :vsn))
result = write_if_changed(path, source)
register_generated_library(library, functions, classes, config, path)
bytes_written = if result == :written, do: byte_size(source), else: 0
SnakeBridge.Telemetry.generate_stop(
start_time,
library.name,
path,
bytes_written,
length(functions),
length(classes)
)
:ok
end
@spec write_if_changed(String.t(), String.t()) :: :written | :unchanged
def write_if_changed(path, new_content) do
with_lock(path, fn ->
case File.read(path) do
{:ok, existing} when existing == new_content ->
:unchanged
_ ->
write_atomic(path, new_content)
end
end)
end
defp write_atomic(path, content) do
temp_path = "#{path}.tmp.#{System.unique_integer([:positive])}"
try do
File.mkdir_p!(Path.dirname(path))
File.write!(temp_path, content)
File.rename!(temp_path, path)
:written
rescue
exception ->
File.rm(temp_path)
reraise exception, __STACKTRACE__
end
end
defp with_lock(path, fun) when is_function(fun, 0) do
lock_key = {:snakebridge_write, Path.expand(path)}
:global.trans(lock_key, fun)
end
defp register_generated_library(library, functions, classes, config, path) do
entry = build_registry_entry(library, functions, classes, config, path)
_ = SnakeBridge.Registry.register(library.python_name, entry)
end
defp build_registry_entry(library, functions, classes, config, path) do
python_version =
case library.version do
nil -> "unknown"
:stdlib -> "stdlib"
value -> to_string(value)
end
%{
python_module: library.python_name,
python_version: python_version,
elixir_module: module_to_string(library.module_name),
generated_at: DateTime.utc_now(),
path: config.generated_dir,
files: [Path.basename(path)],
stats: %{
functions: length(functions),
classes: length(classes),
submodules: count_submodules(library, functions, classes)
}
}
end
defp count_submodules(library, functions, classes) do
base = library.python_name
function_modules =
functions
|> Enum.map(&(&1["python_module"] || &1[:python_module] || base))
class_modules =
classes
|> Enum.map(&(&1["python_module"] || &1[:python_module] || base))
(function_modules ++ class_modules)
|> Enum.uniq()
|> Enum.reject(&(&1 == base))
|> length()
end
@spec render_function(map(), SnakeBridge.Config.Library.t()) :: String.t()
def render_function(info, library) do
raw_name = info["name"] || ""
python_name = info["python_name"] || info["function"] || raw_name
{name, _python_name} = sanitize_function_name(raw_name)
if module_attribute?(info) do
render_module_attribute(name, python_name, info)
else
render_callable_function(info, library, name, python_name)
end
end
defp render_callable_function(info, library, name, python_name) do
params = info["parameters"] || []
doc = info["docstring"] || ""
plan = build_params(params, info)
param_names = Enum.map(plan.required, & &1.name)
args_name = extra_args_name(param_names)
return_type = info["return_type"] || %{"type" => "any"}
normal = render_function_body(name, python_name, plan, args_name, return_type, doc, params)
maybe_add_streaming(normal, name, python_name, plan, args_name, library)
end
defp render_function_body(name, python_name, plan, args_name, return_type, doc, params) do
if plan.is_variadic do
render_variadic_function(name, python_name, return_type, doc, params)
else
render_normal_function(name, python_name, plan, args_name, return_type, doc, params)
end
end
defp maybe_add_streaming(normal, name, python_name, plan, args_name, library) do
is_streaming = python_name in (library.streaming || [])
if is_streaming do
streaming = render_streaming_body(name, python_name, plan, args_name)
normal <> "\n\n" <> streaming
else
normal
end
end
defp render_streaming_body(name, python_name, plan, args_name) do
if plan.is_variadic do
render_variadic_streaming_variant(name, python_name)
else
render_streaming_variant(name, python_name, plan, args_name)
end
end
defp module_attribute?(info) do
info["call_type"] == "module_attr" or info[:call_type] == "module_attr" or
info["type"] == "attribute" or info[:type] == "attribute"
end
defp render_module_attribute(name, python_name, info) do
return_type = info["return_type"] || %{"type" => "any"}
doc = info["docstring"] || ""
formatted_doc = format_docstring(doc, [], return_type)
attr_ref = function_ref(name, python_name)
return_spec = type_spec_string(return_type)
"""
@doc \"\"\"
#{String.trim(formatted_doc)}
\"\"\"
@spec #{name}() :: {:ok, #{return_spec}} | {:error, Snakepit.Error.t()}
def #{name}() do
SnakeBridge.Runtime.get_module_attr(__MODULE__, #{attr_ref})
end
"""
end
defp render_normal_function(name, python_name, plan, args_name, return_type, doc, params) do
param_names = Enum.map(plan.required, & &1.name)
args = args_expr(param_names, plan.has_args, args_name)
call = runtime_call(name, python_name, args)
spec = function_spec(name, plan.required, plan.has_args, return_type)
formatted_doc = format_docstring(doc, params, return_type)
normalize = normalize_args_line(plan.has_args, args_name, 8)
kw_validation = keyword_only_validation(plan.required_keyword_only, 8)
"""
@doc \"\"\"
#{String.trim(formatted_doc)}
\"\"\"
#{spec}
def #{name}(#{param_list(param_names, plan.has_args, plan.has_opts, args_name)}) do
#{normalize}#{kw_validation} #{call}
end
"""
end
defp render_streaming_variant(name, python_name, plan, args_name) do
param_names = Enum.map(plan.required, & &1.name)
args = args_expr(param_names, plan.has_args, args_name)
stream_params =
param_names
|> maybe_add_args(plan.has_args, args_name)
|> Kernel.++(["opts \\\\ []", "callback"])
stream_params_str = Enum.join(stream_params, ", ")
stream_call = runtime_stream_call(name, python_name, args)
spec_args =
plan.required
|> Enum.map(¶m_type_spec/1)
|> maybe_add_args_spec(plan.has_args)
|> Kernel.++(["keyword()", "(term() -> any())"])
spec_args_str = Enum.join(spec_args, ", ")
base_arity = length(param_names) + if(plan.has_args, do: 2, else: 1)
normalize = normalize_args_line(plan.has_args, args_name, 8)
kw_validation = keyword_only_validation(plan.required_keyword_only, 8)
"""
@doc \"\"\"
Streaming variant of `#{name}/#{base_arity}`.
The callback receives chunks as they arrive.
\"\"\"
@spec #{name}_stream(#{spec_args_str}) :: :ok | {:error, Snakepit.Error.t()}
def #{name}_stream(#{stream_params_str}) when is_function(callback, 1) do
#{normalize}#{kw_validation} #{stream_call}
end
"""
end
defp render_variadic_function(name, python_name, return_type, doc, params) do
max_arity = variadic_max_arity()
return_spec = type_spec_string(return_type)
formatted_doc = format_docstring(doc, params, return_type)
specs = variadic_specs(name, max_arity, return_spec)
clauses = variadic_function_clauses(name, python_name, max_arity)
"""
@doc \"\"\"
#{String.trim(formatted_doc)}
\"\"\"
#{specs}
#{indent(clauses, 6)}
"""
end
defp render_variadic_streaming_variant(name, python_name) do
max_arity = variadic_max_arity()
specs = variadic_streaming_specs(name, max_arity)
clauses = variadic_streaming_clauses(name, python_name, max_arity)
"""
@doc \"\"\"
Streaming variant of `#{name}`.
The callback receives chunks as they arrive.
\"\"\"
#{specs}
#{indent(clauses, 6)}
"""
end
defp render_variadic_constructor(_plan, _args_name) do
max_arity = variadic_max_arity()
specs = variadic_specs("new", max_arity, "SnakeBridge.Ref.t()")
clauses = variadic_constructor_clauses(max_arity)
"""
#{specs}
#{indent(clauses, 8)}
"""
end
defp render_variadic_method(name, python_name, return_type) do
max_arity = variadic_max_arity()
return_spec = type_spec_string(return_type)
specs = variadic_method_specs(name, max_arity, return_spec)
clauses = variadic_method_clauses(name, python_name, max_arity)
"""
#{specs}
#{indent(clauses, 8)}
"""
end
defp render_submodule(python_module, functions, library) do
module_name =
python_module
|> String.split(".")
|> Enum.drop(length(String.split(library.python_name, ".")))
|> Enum.map_join(".", &Macro.camelize/1)
function_defs =
functions
|> Enum.sort_by(& &1["name"])
|> Enum.map_join("\n\n", &render_function(&1, library))
"""
defmodule #{module_name} do
def __snakebridge_python_name__, do: "#{python_module}"
def __snakebridge_library__, do: "#{library.python_name}"
#{indent(function_defs, 4)}
end
"""
end
@spec render_class(map(), SnakeBridge.Config.Library.t()) :: String.t()
def render_class(class_info, library) do
class_name = class_name(class_info)
python_module = class_python_module(class_info, library)
module_name = class_module_name(class_info, library)
relative_module = relative_module_name(library, module_name)
methods = class_info["methods"] || []
attrs = class_info["attributes"] || []
init_method = Enum.find(methods, fn method -> method["name"] == "__init__" end)
init_params = if init_method, do: init_method["parameters"] || [], else: []
plan = build_params(init_params, init_method || %{})
param_names = Enum.map(plan.required, & &1.name)
args_name = extra_args_name(param_names)
constructor =
if plan.is_variadic do
render_variadic_constructor(plan, args_name)
else
render_constructor(plan, args_name)
end
methods_source =
methods
|> Enum.reject(fn method -> method["name"] == "__init__" end)
|> Enum.map_join("\n\n", &render_method/1)
attrs_source =
attrs
|> Enum.map_join("\n\n", &render_attribute/1)
"""
defmodule #{relative_module} do
def __snakebridge_python_name__, do: "#{python_module}"
def __snakebridge_python_class__, do: "#{class_name}"
def __snakebridge_library__, do: "#{library.python_name}"
@opaque t :: SnakeBridge.Ref.t()
#{indent(constructor, 4)}
#{indent(methods_source, 4)}
#{indent(attrs_source, 4)}
end
"""
end
defp render_constructor(plan, args_name) do
param_names = Enum.map(plan.required, & &1.name)
args = args_expr(param_names, plan.has_args, args_name)
param_list = param_list(param_names, plan.has_args, plan.has_opts, args_name)
call = "SnakeBridge.Runtime.call_class(__MODULE__, :__init__, #{args}, opts)"
spec_args =
plan.required
|> Enum.map(¶m_type_spec/1)
|> maybe_add_args_spec(plan.has_args)
|> Kernel.++(["keyword()"])
spec_args_str = Enum.join(spec_args, ", ")
normalize = normalize_args_line(plan.has_args, args_name, 10)
kw_validation = keyword_only_validation(plan.required_keyword_only, 10)
"""
@spec new(#{spec_args_str}) :: {:ok, SnakeBridge.Ref.t()} | {:error, Snakepit.Error.t()}
def new(#{param_list}) do
#{normalize}#{kw_validation} #{call}
end
"""
end
defp render_method(%{"name" => "__init__"}), do: ""
defp render_method(%{name: "__init__"}), do: ""
defp render_method(info) do
python_name = info["python_name"] || info["name"] || info[:name] || ""
name_info = resolve_method_name(info, python_name)
do_render_method(name_info, info)
end
defp resolve_method_name(info, python_name) do
case info["elixir_name"] || info[:elixir_name] do
elixir_name when is_binary(elixir_name) -> {elixir_name, python_name}
_ -> sanitize_method_name(python_name)
end
end
defp do_render_method(nil, _info), do: ""
defp do_render_method({name, python_name}, info) do
params = info["parameters"] || []
plan = build_params(params, info)
return_type = info["return_type"] || %{"type" => "any"}
render_method_body(name, python_name, plan, return_type)
end
defp render_method_body(name, python_name, %{is_variadic: true}, return_type) do
render_variadic_method(name, python_name, return_type)
end
defp render_method_body(name, python_name, plan, return_type) do
param_names = Enum.map(plan.required, & &1.name)
args_name = extra_args_name(param_names)
spec = method_spec(name, plan.required, plan.has_args, return_type)
call = runtime_method_call(name, python_name, param_names, plan.has_args, args_name)
normalize = normalize_args_line(plan.has_args, args_name, 10)
kw_validation = keyword_only_validation(plan.required_keyword_only, 10)
"""
#{spec}
def #{name}(ref#{method_param_suffix(param_names, plan.has_args, plan.has_opts, args_name)}) do
#{normalize}#{kw_validation} #{call}
end
"""
end
defp render_attribute(attr) do
"""
@spec #{attr}(SnakeBridge.Ref.t()) :: {:ok, term()} | {:error, Snakepit.Error.t()}
def #{attr}(ref) do
SnakeBridge.Runtime.get_attr(ref, :#{attr})
end
"""
end
defp variadic_max_arity do
Application.get_env(:snakebridge, :variadic_max_arity, 8)
end
defp variadic_specs(name, max_arity, return_spec) do
Enum.map_join(0..max_arity, "\n", fn arity ->
args = variadic_term_args(arity)
variadic_spec_pair(name, args, return_spec)
end)
end
defp variadic_term_args(0), do: []
defp variadic_term_args(arity), do: Enum.map(1..arity, fn _ -> "term()" end)
defp variadic_spec_pair(name, args, return_spec) do
spec_no_opts =
"@spec #{name}(#{Enum.join(args, ", ")}) :: {:ok, #{return_spec}} | {:error, Snakepit.Error.t()}"
spec_with_opts =
"@spec #{name}(#{Enum.join(args ++ ["keyword()"], ", ")}) :: {:ok, #{return_spec}} | {:error, Snakepit.Error.t()}"
spec_no_opts <> "\n" <> spec_with_opts
end
defp variadic_method_specs(name, max_arity, return_spec) do
Enum.map_join(0..max_arity, "\n", fn arity ->
args = ["SnakeBridge.Ref.t()" | variadic_term_args(arity)]
variadic_spec_pair(name, args, return_spec)
end)
end
defp variadic_streaming_specs(name, max_arity) do
Enum.map_join(0..max_arity, "\n", fn arity ->
args = variadic_term_args(arity)
variadic_streaming_spec_pair(name, args)
end)
end
defp variadic_streaming_spec_pair(name, args) do
callback = "(term() -> any())"
spec_no_opts =
"@spec #{name}_stream(#{Enum.join(args ++ [callback], ", ")}) :: :ok | {:error, Snakepit.Error.t()}"
spec_with_opts =
"@spec #{name}_stream(#{Enum.join(args ++ ["keyword()", callback], ", ")}) :: :ok | {:error, Snakepit.Error.t()}"
spec_no_opts <> "\n" <> spec_with_opts
end
defp variadic_function_clauses(name, python_name, max_arity) do
Enum.map_join(0..max_arity, "\n\n", fn arity ->
args = variadic_args(arity)
args_list = variadic_args_list(args)
build_variadic_function_clause(name, python_name, arity, args, args_list)
end)
end
defp build_variadic_function_clause(name, python_name, 0, args, args_list) do
no_args_clause =
variadic_no_opts_clause(name, python_name, variadic_param_list(args), args_list)
opts_clause =
variadic_opts_clause(
name,
python_name,
variadic_param_list_with_opts(args),
args_list
)
no_args_clause <> "\n\n" <> opts_clause
end
defp build_variadic_function_clause(name, python_name, _arity, args, args_list) do
positional_clause =
variadic_no_opts_clause(name, python_name, variadic_param_list(args), args_list)
opts_clause =
variadic_opts_clause(
name,
python_name,
variadic_param_list_with_opts(args),
args_list
)
positional_clause <> "\n\n" <> opts_clause
end
defp variadic_constructor_clauses(max_arity) do
Enum.map_join(0..max_arity, "\n\n", fn arity ->
args = variadic_args(arity)
args_list = variadic_args_list(args)
build_variadic_constructor_clause(arity, args, args_list)
end)
end
defp build_variadic_constructor_clause(0, args, args_list) do
no_args_clause =
variadic_constructor_no_opts_clause(variadic_param_list(args), args_list)
opts_clause =
variadic_constructor_opts_clause(
variadic_param_list_with_opts(args),
args_list
)
no_args_clause <> "\n\n" <> opts_clause
end
defp build_variadic_constructor_clause(_arity, args, args_list) do
positional_clause =
variadic_constructor_no_opts_clause(variadic_param_list(args), args_list)
opts_clause =
variadic_constructor_opts_clause(
variadic_param_list_with_opts(args),
args_list
)
positional_clause <> "\n\n" <> opts_clause
end
defp variadic_method_clauses(name, python_name, max_arity) do
Enum.map_join(0..max_arity, "\n\n", fn arity ->
args = variadic_args(arity)
args_list = variadic_args_list(args)
build_variadic_method_clause(name, python_name, arity, args, args_list)
end)
end
defp build_variadic_method_clause(name, python_name, 0, args, args_list) do
no_args_clause =
variadic_method_no_opts_clause(
name,
python_name,
variadic_method_param_list(args),
args_list
)
opts_clause =
variadic_method_opts_clause(
name,
python_name,
variadic_method_param_list_with_opts(args),
args_list
)
no_args_clause <> "\n\n" <> opts_clause
end
defp build_variadic_method_clause(name, python_name, _arity, args, args_list) do
positional_clause =
variadic_method_no_opts_clause(
name,
python_name,
variadic_method_param_list(args),
args_list
)
opts_clause =
variadic_method_opts_clause(
name,
python_name,
variadic_method_param_list_with_opts(args),
args_list
)
positional_clause <> "\n\n" <> opts_clause
end
defp variadic_streaming_clauses(name, python_name, max_arity) do
Enum.map_join(0..max_arity, "\n\n", fn arity ->
args = variadic_args(arity)
args_list = variadic_args_list(args)
build_variadic_streaming_clause(name, python_name, arity, args, args_list)
end)
end
defp build_variadic_streaming_clause(name, python_name, 0, args, args_list) do
no_args_clause =
variadic_streaming_no_opts_clause(
name,
python_name,
variadic_streaming_param_list(args),
args_list
)
opts_clause =
variadic_streaming_opts_clause(
name,
python_name,
variadic_streaming_param_list_with_opts(args),
args_list
)
no_args_clause <> "\n\n" <> opts_clause
end
defp build_variadic_streaming_clause(name, python_name, _arity, args, args_list) do
positional_clause =
variadic_streaming_no_opts_clause(
name,
python_name,
variadic_streaming_param_list(args),
args_list
)
opts_clause =
variadic_streaming_opts_clause(
name,
python_name,
variadic_streaming_param_list_with_opts(args),
args_list
)
positional_clause <> "\n\n" <> opts_clause
end
defp variadic_no_opts_clause(name, python_name, params, args_list) do
call = runtime_call(name, python_name, args_list, "[]")
"""
def #{name}(#{params}) do
#{call}
end
"""
end
defp variadic_opts_clause(name, python_name, params, args_list) do
call = runtime_call(name, python_name, args_list, "opts")
"""
def #{name}(#{params}) when #{opts_guard()} do
#{call}
end
"""
end
defp variadic_constructor_no_opts_clause(params, args_list) do
call = "SnakeBridge.Runtime.call_class(__MODULE__, :__init__, #{args_list}, [])"
"""
def new(#{params}) do
#{call}
end
"""
end
defp variadic_constructor_opts_clause(params, args_list) do
call = "SnakeBridge.Runtime.call_class(__MODULE__, :__init__, #{args_list}, opts)"
"""
def new(#{params}) when #{opts_guard()} do
#{call}
end
"""
end
defp variadic_method_no_opts_clause(name, python_name, params, args_list) do
call =
"SnakeBridge.Runtime.call_method(ref, #{function_ref(name, python_name)}, #{args_list}, [])"
"""
def #{name}(#{params}) do
#{call}
end
"""
end
defp variadic_method_opts_clause(name, python_name, params, args_list) do
call =
"SnakeBridge.Runtime.call_method(ref, #{function_ref(name, python_name)}, #{args_list}, opts)"
"""
def #{name}(#{params}) when #{opts_guard()} do
#{call}
end
"""
end
defp variadic_streaming_no_opts_clause(name, python_name, params, args_list) do
call = runtime_stream_call(name, python_name, args_list, "[]")
"""
def #{name}_stream(#{params}) when is_function(callback, 1) do
#{call}
end
"""
end
defp variadic_streaming_opts_clause(name, python_name, params, args_list) do
call = runtime_stream_call(name, python_name, args_list, "opts")
"""
def #{name}_stream(#{params}) when #{opts_guard()} and is_function(callback, 1) do
#{call}
end
"""
end
defp variadic_args(arity) when is_integer(arity) and arity > 0 do
Enum.map(1..arity, &"arg#{&1}")
end
defp variadic_args(_arity), do: []
defp variadic_args_list([]), do: "[]"
defp variadic_args_list(args), do: "[" <> Enum.join(args, ", ") <> "]"
defp variadic_param_list([]), do: ""
defp variadic_param_list(args), do: Enum.join(args, ", ")
defp variadic_param_list_with_opts([]), do: "opts"
defp variadic_param_list_with_opts(args), do: Enum.join(args ++ ["opts"], ", ")
defp variadic_method_param_list(args), do: Enum.join(["ref" | args], ", ")
defp variadic_method_param_list_with_opts(args), do: Enum.join(["ref" | args] ++ ["opts"], ", ")
defp variadic_streaming_param_list(args), do: Enum.join(args ++ ["callback"], ", ")
defp variadic_streaming_param_list_with_opts(args),
do: Enum.join(args ++ ["opts", "callback"], ", ")
defp opts_guard do
"is_list(opts) and (opts == [] or (is_tuple(hd(opts)) and tuple_size(hd(opts)) == 2 and is_atom(elem(hd(opts), 0))))"
end
defp render_discovery(functions, classes) do
function_list =
functions
|> Enum.map_join(",\n ", fn info ->
name =
case info["elixir_name"] do
elixir when is_binary(elixir) ->
elixir
_ ->
info["name"] |> to_string() |> sanitize_function_name() |> elem(0)
end
arity = required_arity(info["parameters"] || [])
summary = info["docstring"] |> to_string() |> String.split("\n") |> List.first() || ""
"{:#{name}, #{arity}, __MODULE__, #{inspect(summary)}}"
end)
class_list =
classes
|> Enum.map_join(",\n ", fn info ->
module = class_module_name(info, nil)
doc = info["docstring"] |> to_string()
"{#{module}, #{inspect(doc)}}"
end)
"""
@doc false
def __functions__ do
[
#{function_list}
]
end
@doc false
def __classes__ do
[
#{class_list}
]
end
@doc false
def __search__(query) do
SnakeBridge.Docs.search(__MODULE__, query)
end
@doc false
def doc(function) do
SnakeBridge.Docs.get(__MODULE__, function)
end
"""
end
@spec format_docstring(String.t() | nil, list(), map() | nil) :: String.t()
def format_docstring(raw_doc, params \\ [], return_type \\ nil)
def format_docstring(nil, _params, _return_type), do: ""
def format_docstring("", _params, _return_type), do: ""
def format_docstring(raw_doc, params, return_type) when is_binary(raw_doc) do
base =
raw_doc
|> RstParser.parse()
|> MarkdownConverter.convert()
extras = format_param_docs(params, return_type)
if extras == "" do
base
else
base <> "\n\n" <> extras
end
rescue
_ ->
extras = format_param_docs(params, return_type)
if extras == "" do
raw_doc
else
raw_doc <> "\n\n" <> extras
end
end
@spec build_params(list(), map()) :: %{
required: list(map()),
has_args: boolean(),
has_opts: boolean(),
is_variadic: boolean(),
required_keyword_only: list(map()),
optional_keyword_only: list(map()),
has_var_keyword: boolean()
}
def build_params(params, info \\ %{}) when is_list(params) do
signature_available = Map.get(info, "signature_available", true)
if params == [] and not signature_available do
%{
required: [],
has_args: true,
has_opts: true,
is_variadic: true,
required_keyword_only: [],
optional_keyword_only: [],
has_var_keyword: false
}
else
required =
params
|> Enum.filter(&required_positional?/1)
|> Enum.map(¶m_entry/1)
required_kw_only = Enum.filter(params, &keyword_only_required?/1)
optional_kw_only = Enum.filter(params, &keyword_only_optional?/1)
has_args =
Enum.any?(params, fn param ->
optional_positional?(param) or varargs?(param)
end)
%{
required: required,
has_args: has_args,
has_opts: true,
is_variadic: false,
required_keyword_only: required_kw_only,
optional_keyword_only: optional_kw_only,
has_var_keyword: Enum.any?(params, &kwargs?/1)
}
end
end
defp required_arity(params) do
params
|> Enum.filter(fn param ->
param_kind(param) in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"]
end)
|> Enum.reject(¶m_default?/1)
|> length()
end
defp param_list(param_names, has_args, has_opts, args_name) do
param_names
|> maybe_add_args(has_args, args_name)
|> maybe_add_opts(has_opts)
|> Enum.join(", ")
end
defp runtime_call(name, python_name, args, opts_expr \\ "opts") do
"SnakeBridge.Runtime.call(__MODULE__, #{function_ref(name, python_name)}, #{args}, #{opts_expr})"
end
defp runtime_stream_call(name, python_name, args, opts_expr \\ "opts") do
"SnakeBridge.Runtime.stream(__MODULE__, #{function_ref(name, python_name)}, #{args}, #{opts_expr}, callback)"
end
defp function_ref(name, python_name) do
if python_name == name do
":#{name}"
else
inspect(python_name)
end
end
defp function_spec(name, param_entries, has_args, return_type) do
args =
param_entries
|> Enum.map(¶m_type_spec/1)
|> maybe_add_args_spec(has_args)
|> Kernel.++(["keyword()"])
return_spec = type_spec_string(return_type)
"@spec #{name}(#{Enum.join(args, ", ")}) :: {:ok, #{return_spec}} | {:error, Snakepit.Error.t()}"
end
defp method_spec(name, param_entries, has_args, return_type) do
args =
[ref_type_spec()]
|> Kernel.++(Enum.map(param_entries, ¶m_type_spec/1))
|> maybe_add_args_spec(has_args)
|> Kernel.++(["keyword()"])
return_spec = type_spec_string(return_type)
"@spec #{name}(#{Enum.join(args, ", ")}) :: {:ok, #{return_spec}} | {:error, Snakepit.Error.t()}"
end
defp runtime_method_call(name, python_name, param_names, has_args, args_name) do
args = args_expr(param_names, has_args, args_name)
"SnakeBridge.Runtime.call_method(ref, #{function_ref(name, python_name)}, #{args}, opts)"
end
defp method_param_suffix(param_names, has_args, has_opts, args_name) do
suffix =
param_names
|> maybe_add_args(has_args, args_name)
|> maybe_add_opts(has_opts)
", " <> Enum.join(suffix, ", ")
end
defp args_expr(param_names, true, args_name) do
base = "[" <> Enum.join(param_names, ", ") <> "]"
base <> " ++ List.wrap(" <> args_name <> ")"
end
defp args_expr(param_names, false, _args_name) do
"[" <> Enum.join(param_names, ", ") <> "]"
end
defp maybe_add_args(items, true, args_name), do: items ++ ["#{args_name} \\\\ []"]
defp maybe_add_args(items, false, _args_name), do: items
defp maybe_add_opts(items, _has_opts), do: items ++ ["opts \\\\ []"]
defp maybe_add_args_spec(items, true), do: items ++ ["list(term())"]
defp maybe_add_args_spec(items, false), do: items
defp normalize_args_line(true, args_name, indent) do
String.duplicate(" ", indent) <>
"{#{args_name}, opts} = SnakeBridge.Runtime.normalize_args_opts(#{args_name}, opts)\n"
end
defp normalize_args_line(false, _args_name, _indent), do: ""
defp keyword_only_validation(required_keyword_only, indent) do
names =
required_keyword_only
|> Enum.map(¶m_name/1)
|> Enum.reject(&is_nil/1)
if names == [] do
""
else
padding = String.duplicate(" ", indent)
"""
#{padding}kw_keys = opts |> Keyword.keys() |> Enum.map(&to_string/1)
#{padding}missing_kw = #{inspect(names)} |> Enum.reject(&(&1 in kw_keys))
#{padding}if missing_kw != [] do
#{padding} raise ArgumentError,
#{padding} "Missing required keyword-only arguments: " <> Enum.join(missing_kw, ", ")
#{padding}end
"""
end
end
defp extra_args_name(param_names) do
if "args" in param_names do
"extra_args"
else
"args"
end
end
defp param_entry(param) do
%{
name: sanitize_name(param),
type: param_type(param)
}
end
defp param_type(%{"type" => type}) when is_map(type), do: type
defp param_type(%{type: type}) when is_map(type), do: type
defp param_type(_), do: %{"type" => "any"}
defp param_type_spec(%{type: type}), do: type_spec_string(type)
defp param_type_spec(_), do: "term()"
defp type_spec_string(type) do
type
|> TypeMapper.to_spec()
|> Macro.to_string()
end
defp ref_type_spec do
"SnakeBridge.Ref.t()"
end
defp required_positional?(param) do
kind = param_kind(param)
kind in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"] and not param_default?(param)
end
defp optional_positional?(param) do
kind = param_kind(param)
kind in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"] and param_default?(param)
end
defp varargs?(param), do: param_kind(param) == "VAR_POSITIONAL"
defp kwargs?(param), do: param_kind(param) == "VAR_KEYWORD"
defp keyword_only_required?(param) do
param_kind(param) == "KEYWORD_ONLY" and not param_default?(param)
end
defp keyword_only_optional?(param) do
param_kind(param) == "KEYWORD_ONLY" and param_default?(param)
end
defp param_kind(%{"kind" => kind}) when is_binary(kind), do: String.upcase(kind)
defp param_kind(%{kind: kind}) when is_binary(kind), do: String.upcase(kind)
defp param_kind(%{kind: kind}), do: kind
defp param_kind(_), do: nil
defp param_default?(%{"default" => _}), do: true
defp param_default?(%{default: _}), do: true
defp param_default?(_), do: false
defp sanitize_name(%{"name" => name}), do: sanitize_name(name)
defp sanitize_name(name) when is_binary(name) do
name
|> Macro.underscore()
|> String.replace(~r/[^a-z0-9_]/, "_")
|> ensure_identifier()
end
defp sanitize_function_name(python_name) when is_binary(python_name) do
elixir_name =
python_name
|> Macro.underscore()
|> String.replace(~r/[^a-z0-9_?!]/, "_")
|> ensure_valid_function_identifier()
elixir_name =
if elixir_name in @reserved_words do
"py_#{elixir_name}"
else
elixir_name
end
{elixir_name, python_name}
end
@doc false
@spec sanitize_method_name(String.t()) :: {String.t(), String.t()} | nil
def sanitize_method_name(python_name) when is_binary(python_name) do
cond do
Map.has_key?(@dunder_mappings, python_name) ->
{Map.get(@dunder_mappings, python_name), python_name}
String.starts_with?(python_name, "__") and String.ends_with?(python_name, "__") ->
nil
python_name in @reserved_words ->
{"py_#{python_name}", python_name}
true ->
elixir_name =
python_name
|> Macro.underscore()
|> String.replace(~r/[^a-z0-9_?!]/, "_")
|> ensure_valid_function_identifier()
{elixir_name, python_name}
end
end
defp ensure_valid_function_identifier(""), do: "_"
defp ensure_valid_function_identifier(name) do
if String.match?(name, ~r/^[a-z_][a-z0-9_?!]*$/) do
name
else
"_" <> name
end
end
defp ensure_identifier(<> = name)
when first in ?a..?z or first == ?_ do
name
end
defp ensure_identifier(<> = name) when first in ?0..?9 do
"_" <> name
end
defp ensure_identifier(""), do: "_unnamed"
defp ensure_identifier(name), do: "_" <> name
defp module_to_string(module) when is_atom(module),
do: module |> Module.split() |> Enum.join(".")
defp module_to_string(module) when is_binary(module), do: module
defp class_sort_key(info) do
{class_module_name(info, nil), class_name(info)}
end
defp class_name(info) do
info["name"] || info["class"] || "Class"
end
defp class_python_module(info, library) do
info["python_module"] || library.python_name
end
defp class_module_name(info, nil) do
info["module"] || class_name(info)
end
defp class_module_name(info, library) do
case info["module"] do
module when is_binary(module) ->
module
_ ->
python_module = class_python_module(info, library)
python_parts = String.split(python_module, ".")
library_parts = String.split(library.python_name, ".")
extra_parts = Enum.drop(python_parts, length(library_parts))
extra_parts = drop_class_suffix(extra_parts, class_name(info))
library.module_name
|> Module.split()
|> Kernel.++(Enum.map(extra_parts, &Macro.camelize/1))
|> Kernel.++([class_name(info)])
|> Module.concat()
|> module_to_string()
end
end
defp relative_module_name(library, module_name) when is_binary(module_name) do
base = module_to_string(library.module_name) <> "."
if String.starts_with?(module_name, base) do
String.replace_prefix(module_name, base, "")
else
module_name
end
end
defp drop_class_suffix(parts, class_name) when is_list(parts) and is_binary(class_name) do
class_suffix = Macro.underscore(class_name)
case List.last(parts) do
^class_suffix -> Enum.drop(parts, -1)
_ -> parts
end
end
defp drop_class_suffix(parts, _class_name), do: parts
defp indent(text, spaces) do
prefix = String.duplicate(" ", spaces)
text
|> String.split("\n")
|> Enum.map_join("\n", fn line -> prefix <> line end)
end
defp format_param_docs(params, return_type) do
param_lines =
params
|> Enum.map(¶m_doc_line/1)
|> Enum.reject(&is_nil/1)
sections =
[]
|> maybe_add_params_section(param_lines)
|> maybe_add_return_section(return_type)
Enum.join(sections, "\n\n")
end
defp maybe_add_params_section(sections, []), do: sections
defp maybe_add_params_section(sections, lines) do
sections ++ ["Parameters:\n" <> Enum.join(lines, "\n")]
end
defp maybe_add_return_section(sections, nil), do: sections
defp maybe_add_return_section(sections, return_type) do
return_spec = type_spec_string(return_type)
sections ++ ["Returns:\n- `#{return_spec}`"]
end
defp param_doc_line(param) do
case param_name(param) do
nil -> nil
name -> format_param_doc_line(name, param)
end
end
defp format_param_doc_line(name, param) do
type = param_type(param) |> type_spec_string()
kind_fragment = param_kind_fragment(param)
default_fragment = param_default_fragment(param)
"- `#{name}` (#{type}#{kind_fragment}#{default_fragment})"
end
defp param_kind_fragment(param) do
case param_kind(param) do
"KEYWORD_ONLY" -> keyword_only_fragment(param)
_ -> ""
end
end
defp keyword_only_fragment(param) do
if param_default?(param), do: " keyword-only", else: " keyword-only, required"
end
defp param_default_fragment(param) do
case param_default_value(param) do
nil -> ""
default -> " default: #{default}"
end
end
defp param_name(%{"name" => name}) when is_binary(name), do: name
defp param_name(%{name: name}) when is_binary(name), do: name
defp param_name(_), do: nil
defp param_default_value(%{"default" => default}), do: default
defp param_default_value(%{default: default}), do: default
defp param_default_value(_), do: nil
end
defmodule SnakeBridge.HelperGenerator do
@moduledoc """
Generates Elixir helper wrappers from registry data.
"""
@spec render_library(String.t(), list(), keyword()) :: String.t()
def render_library(library, helpers, opts \\ []) when is_binary(library) do
version = Keyword.get(opts, :version, Application.spec(:snakebridge, :vsn) |> to_string())
header = """
# Generated by SnakeBridge v#{version} - DO NOT EDIT MANUALLY
# Regenerate with: mix compile
# Helper library: #{library}
"""
modules = render_modules(library, helpers)
"""
#{header}
#{modules}
"""
|> Code.format_string!()
|> IO.iodata_to_binary()
end
@spec generate_helpers(list(), SnakeBridge.Config.t()) :: :ok
def generate_helpers(helpers, config) do
grouped = group_by_library(helpers)
if map_size(grouped) == 0 do
:ok
else
dir = Path.join(config.generated_dir, "helpers")
File.mkdir_p!(dir)
Enum.each(grouped, fn {library, entries} ->
source = render_library(library, entries, version: Application.spec(:snakebridge, :vsn))
path = Path.join(dir, "#{library}.ex")
SnakeBridge.Generator.write_if_changed(path, source)
end)
:ok
end
end
defp group_by_library(helpers) do
helpers
|> Enum.map(&helper_name/1)
|> Enum.zip(helpers)
|> Enum.reduce(%{}, fn
{nil, _}, acc ->
acc
{name, helper}, acc ->
case String.split(name, ".", parts: 2) do
[library, _rest] ->
Map.update(acc, library, [helper], fn existing -> [helper | existing] end)
_ ->
acc
end
end)
|> Map.new(fn {library, entries} ->
{library, Enum.reverse(entries)}
end)
end
defp helper_name(%{"name" => name}) when is_binary(name), do: name
defp helper_name(%{name: name}) when is_binary(name), do: name
defp helper_name(_), do: nil
defp render_modules(library, helpers) do
root = root_module(library)
helpers
|> Enum.map(&normalize_helper(library, &1))
|> Enum.group_by(& &1.module)
|> Enum.sort_by(fn {module, _} -> module end)
|> Enum.map_join("\n\n", fn {module, entries} ->
moduledoc = if module == root, do: root_moduledoc(library), else: " @moduledoc false"
functions =
entries |> Enum.sort_by(& &1.function) |> Enum.map_join("\n\n", &render_function/1)
"""
defmodule #{module} do
#{moduledoc}
#{indent(functions, 2)}
end
"""
end)
end
defp root_module(library) do
[Macro.camelize(library), "Helpers"] |> Enum.join(".")
end
defp root_moduledoc(library) do
" @moduledoc #{inspect("Helper wrappers for `#{library}`.")}"
end
defp normalize_helper(library, helper) do
helper_name = helper_name(helper) || ""
segments = String.split(helper_name, ".")
module_segments =
segments
|> Enum.drop(1)
|> Enum.drop(-1)
|> Enum.map(&Macro.camelize/1)
module = ([Macro.camelize(library), "Helpers"] ++ module_segments) |> Enum.join(".")
%{
helper: helper_name,
module: module,
function: helper_function_name(segments),
parameters: helper_parameters(helper),
docstring: helper_docstring(helper)
}
end
defp helper_function_name(segments) do
segments
|> List.last()
|> sanitize_name()
end
defp helper_parameters(%{"parameters" => params}) when is_list(params), do: params
defp helper_parameters(%{parameters: params}) when is_list(params), do: params
defp helper_parameters(_), do: []
defp helper_docstring(%{"docstring" => doc}) when is_binary(doc), do: doc
defp helper_docstring(%{docstring: doc}) when is_binary(doc), do: doc
defp helper_docstring(_), do: ""
defp render_function(%{helper: helper_name, function: name, parameters: params, docstring: doc}) do
plan = SnakeBridge.Generator.build_params(params)
param_names = Enum.map(plan.required, & &1.name)
args_name = extra_args_name(param_names)
args = args_expr(param_names, plan.has_args, args_name)
call = helper_call(helper_name, args)
spec = function_spec(name, plan.required, plan.has_args)
normalize = normalize_args_line(plan.has_args, args_name, 8)
docstring = String.trim(doc)
"""
@doc #{inspect(docstring)}
#{spec}
def #{name}(#{param_list(param_names, plan.has_args, plan.has_opts, args_name)}) do
#{normalize} #{call}
end
"""
end
defp helper_call(name, args) do
"SnakeBridge.Runtime.call_helper(\"#{name}\", #{args}, opts)"
end
defp function_spec(name, param_entries, has_args) do
args = Enum.map(param_entries, fn _ -> "term()" end)
args = if has_args, do: args ++ ["list(term())"], else: args
args = args ++ ["keyword()"]
"@spec #{name}(#{Enum.join(args, ", ")}) :: {:ok, term()} | {:error, term()}"
end
defp param_list(param_names, has_args, has_opts, args_name) do
param_names
|> maybe_add_args(has_args, args_name)
|> maybe_add_opts(has_opts)
|> Enum.join(", ")
end
defp args_expr(param_names, true, args_name) do
base = "[" <> Enum.join(param_names, ", ") <> "]"
base <> " ++ List.wrap(" <> args_name <> ")"
end
defp args_expr(param_names, false, _args_name) do
"[" <> Enum.join(param_names, ", ") <> "]"
end
defp maybe_add_args(items, true, args_name), do: items ++ ["#{args_name} \\\\ []"]
defp maybe_add_args(items, false, _args_name), do: items
defp maybe_add_opts(items, _has_opts), do: items ++ ["opts \\\\ []"]
defp normalize_args_line(true, args_name, indent) do
String.duplicate(" ", indent) <>
"{#{args_name}, opts} = SnakeBridge.Runtime.normalize_args_opts(#{args_name}, opts)\n"
end
defp normalize_args_line(false, _args_name, _indent), do: ""
defp extra_args_name(param_names) do
if "args" in param_names do
"extra_args"
else
"args"
end
end
defp sanitize_name(%{"name" => name}), do: sanitize_name(name)
defp sanitize_name(name) when is_binary(name) do
name
|> Macro.underscore()
|> String.replace(~r/[^a-z0-9_]/, "_")
|> ensure_identifier()
end
defp ensure_identifier(<> = name)
when first in ?a..?z or first == ?_ do
name
end
defp ensure_identifier(<> = name) when first in ?0..?9 do
"_" <> name
end
defp ensure_identifier(""), do: "_unnamed"
defp ensure_identifier(name), do: "_" <> name
defp indent(text, spaces) do
prefix = String.duplicate(" ", spaces)
text
|> String.split("\n")
|> Enum.map_join("\n", fn line -> prefix <> line end)
end
end
defmodule SnakeBridge.HelperNotFoundError do
@moduledoc """
Error raised when a helper name is not registered.
"""
defexception [:message, :helper, :suggestion]
@type t :: %__MODULE__{
message: String.t(),
helper: String.t() | nil,
suggestion: String.t() | nil
}
@impl Exception
def message(%__MODULE__{message: message, suggestion: suggestion}) do
if suggestion do
message <> "\n\nSuggestion: " <> suggestion
else
message
end
end
@spec new(String.t()) :: t()
def new(helper) do
%__MODULE__{
helper: helper,
message: "Helper '#{helper}' not found",
suggestion: "Add a helper under priv/python/helpers or enable the helper pack"
}
end
end
defmodule SnakeBridge.HelperRegistryError do
@moduledoc """
Error raised when helper registry discovery fails.
"""
defexception [:type, :message, :python_error, :suggestion]
@type t :: %__MODULE__{
type: :load_failed,
message: String.t(),
python_error: String.t() | nil,
suggestion: String.t() | nil
}
@impl Exception
def message(%__MODULE__{message: message, suggestion: suggestion}) do
if suggestion do
message <> "\n\nSuggestion: " <> suggestion
else
message
end
end
@doc """
Build an error from Python stderr output.
"""
@spec from_python_output(String.t()) :: t()
def from_python_output(output) when is_binary(output) do
%__MODULE__{
type: :load_failed,
message: "Helper registry failed to load",
python_error: String.trim(output),
suggestion: "Check helper paths or disable the helper pack"
}
end
end
defmodule SnakeBridge.Helpers do
@moduledoc """
Helper registry discovery and configuration for SnakeBridge.
"""
alias SnakeBridge.{Config, HelperRegistryError}
@type helper_info :: map()
@spec discover() :: {:ok, [helper_info()]} | {:error, term()}
def discover do
discover(runtime_config())
end
@spec discover(Config.t() | map()) :: {:ok, [helper_info()]} | {:error, term()}
def discover(%Config{} = config) do
discover(config_to_map(config))
end
def discover(%{} = config) do
payload = payload_config(config, include_adapter_root: true)
case python_runner().run(helper_index_script(), [Jason.encode!(payload)], runner_opts()) do
{:ok, output} ->
parse_output(output)
{:error, {:python_exit, _status, output}} ->
{:error, HelperRegistryError.from_python_output(output)}
{:error, reason} ->
{:error, reason}
end
end
@spec runtime_config() :: map()
def runtime_config do
%{
helper_paths: Application.get_env(:snakebridge, :helper_paths, ["priv/python/helpers"]),
helper_pack_enabled: Application.get_env(:snakebridge, :helper_pack_enabled, true),
helper_allowlist: Application.get_env(:snakebridge, :helper_allowlist, :all),
inline_enabled: Application.get_env(:snakebridge, :inline_enabled, false)
}
end
@spec enabled?(map()) :: boolean()
def enabled?(%{} = config) do
normalized = normalize_config(config)
normalized.helper_allowlist != [] and
(normalized.helper_pack_enabled == true or normalized.helper_paths != [])
end
@spec payload_config(map(), keyword()) :: map()
def payload_config(%{} = config, opts \\ []) do
normalized = normalize_config(config)
payload = %{
"helper_paths" => normalized.helper_paths,
"helper_pack_enabled" => normalized.helper_pack_enabled,
"helper_allowlist" => allowlist_payload(normalized.helper_allowlist)
}
if Keyword.get(opts, :include_adapter_root, false) do
Map.put(payload, "adapter_root", adapter_root())
else
payload
end
end
defp config_to_map(%Config{} = config) do
%{
helper_paths: config.helper_paths,
helper_pack_enabled: config.helper_pack_enabled,
helper_allowlist: config.helper_allowlist,
inline_enabled: config.inline_enabled
}
end
defp normalize_config(%{} = config) do
%{
helper_paths: normalize_paths(Map.get(config, :helper_paths, ["priv/python/helpers"])),
helper_pack_enabled: Map.get(config, :helper_pack_enabled, true) == true,
helper_allowlist: normalize_allowlist(Map.get(config, :helper_allowlist, :all)),
inline_enabled: Map.get(config, :inline_enabled, false) == true
}
end
defp normalize_paths(paths) do
paths
|> List.wrap()
|> Enum.map(&to_string/1)
|> Enum.map(&Path.expand/1)
|> Enum.uniq()
end
defp normalize_allowlist(:all), do: :all
defp normalize_allowlist("all"), do: :all
defp normalize_allowlist(nil), do: :all
defp normalize_allowlist(:none), do: []
defp normalize_allowlist(list) when is_list(list) do
list
|> Enum.map(&to_string/1)
|> Enum.uniq()
end
defp normalize_allowlist(other), do: [to_string(other)]
defp allowlist_payload(:all), do: "all"
defp allowlist_payload(list) when is_list(list), do: list
defp adapter_root do
:snakebridge
|> :code.priv_dir()
|> to_string()
|> Path.join("python")
end
defp python_runner do
Application.get_env(:snakebridge, :python_runner, SnakeBridge.PythonRunner.System)
end
defp runner_opts do
config = Application.get_env(:snakebridge, :helper_registry, [])
Keyword.take(config, [:timeout, :env, :cd])
end
defp parse_output(output) do
case Jason.decode(output) do
{:ok, results} when is_list(results) -> {:ok, results}
{:ok, %{"error" => error}} -> {:error, error}
{:error, _} -> {:error, {:json_parse, output}}
end
end
defp helper_index_script do
~S"""
import json
import sys
config = json.loads(sys.argv[1])
adapter_root = config.pop("adapter_root", None)
if adapter_root:
sys.path.insert(0, adapter_root)
from snakebridge_adapter import helper_registry_index
index = helper_registry_index(config)
print(json.dumps(index))
"""
end
end
defmodule SnakeBridge.IntrospectionError do
@moduledoc """
Structured error for Python introspection failures.
"""
defexception [:type, :package, :message, :python_error, :suggestion]
@type t :: %__MODULE__{
type: :package_not_found | :import_error | :timeout | :introspection_bug,
package: String.t() | nil,
message: String.t(),
python_error: String.t() | nil,
suggestion: String.t() | nil
}
@impl Exception
def message(%__MODULE__{message: message, suggestion: suggestion}) do
if suggestion do
message <> "\n\nSuggestion: " <> suggestion
else
message
end
end
@doc """
Parses Python stderr to classify the error.
"""
@spec from_python_output(String.t(), String.t()) :: t()
def from_python_output(output, package) when is_binary(output) do
output = String.trim(output)
cond do
module_not_found?(output) ->
missing = extract_missing_package(output) || package
%__MODULE__{
type: :package_not_found,
package: missing,
message: "Package '#{missing}' not found",
python_error: output,
suggestion: "Run: mix snakebridge.setup"
}
import_error?(output) ->
%__MODULE__{
type: :import_error,
package: package,
message: extract_import_error(output),
python_error: output,
suggestion: "Check library dependencies or install optional extras"
}
timeout_error?(output) ->
%__MODULE__{
type: :timeout,
package: package,
message: extract_timeout_error(output),
python_error: output,
suggestion: "Increase introspection timeout or retry"
}
true ->
%__MODULE__{
type: :introspection_bug,
package: package,
message: extract_generic_error(output),
python_error: output,
suggestion: "Please report this issue with the Python error output"
}
end
end
defp module_not_found?(output) do
String.contains?(output, "ModuleNotFoundError")
end
defp import_error?(output) do
String.contains?(output, "ImportError")
end
defp timeout_error?(output) do
String.contains?(output, "TimeoutError") or String.contains?(output, "timed out")
end
defp extract_missing_package(output) do
case Regex.run(~r/ModuleNotFoundError: No module named ['"]([^'"]+)['"]/m, output) do
[_, name] -> name
_ -> nil
end
end
defp extract_import_error(output) do
case Regex.run(~r/ImportError: (.+)$/m, output) do
[_, msg] -> String.trim(msg)
_ -> "ImportError"
end
end
defp extract_timeout_error(output) do
case Regex.run(~r/TimeoutError: (.+)$/m, output) do
[_, msg] -> String.trim(msg)
_ -> "Introspection timed out"
end
end
defp extract_generic_error(output) do
output
|> String.split("\n", trim: true)
|> Enum.reverse()
|> Enum.find(&String.contains?(&1, "Error"))
|> case do
nil -> "Unexpected error during introspection"
line -> String.trim(line)
end
end
end
defmodule SnakeBridge.Introspector do
@moduledoc """
Introspects Python functions using the standalone introspection script.
"""
alias SnakeBridge.IntrospectionError
@type function_name :: atom() | String.t()
@spec introspect(SnakeBridge.Config.Library.t() | map(), [function_name()]) ::
{:ok, map()} | {:error, term()}
def introspect(library, functions) when is_list(functions) do
introspect(library, functions, nil)
end
@spec introspect(SnakeBridge.Config.Library.t() | map(), [function_name()], String.t() | nil) ::
{:ok, map()} | {:error, term()}
def introspect(library, functions, python_module) when is_list(functions) do
case introspect_symbols(library, functions, python_module) do
{:ok, infos} ->
{:ok, group_symbols(infos)}
{:error, _reason} = error ->
error
end
end
@spec introspect_symbols(
SnakeBridge.Config.Library.t() | map(),
[function_name()],
String.t() | nil
) :: {:ok, list()} | {:error, term()}
defp introspect_symbols(library, functions, python_module) when is_list(functions) do
start_time = System.monotonic_time()
library_label = library_label(library)
SnakeBridge.Telemetry.introspect_start(library_label, length(functions))
python_name = python_module || library_python_name(library)
functions_json = Jason.encode!(Enum.map(functions, &to_string/1))
result =
case run_script(
[
script_path(),
"--module",
python_name,
"--symbols",
functions_json
],
runner_opts()
) do
{output, 0} -> parse_output(output)
{output, _status} -> handle_python_error(output, python_name)
end
symbols =
case result do
{:ok, results} when is_list(results) -> length(results)
_ -> 0
end
SnakeBridge.Telemetry.introspect_stop(
start_time,
library_label,
symbols,
0,
System.monotonic_time() - start_time
)
result
end
@spec introspect_batch([
{SnakeBridge.Config.Library.t() | map(), String.t(), [function_name()]}
]) ::
list(
{SnakeBridge.Config.Library.t() | map(), {:ok, list()} | {:error, term()}, String.t()}
)
def introspect_batch(libs_and_functions) when is_list(libs_and_functions) do
nested_config = Application.get_env(:snakebridge, :introspector, [])
default_timeout = Application.get_env(:snakebridge, :introspector_timeout, 30_000)
default_concurrency =
Application.get_env(:snakebridge, :introspector_max_concurrency, System.schedulers_online())
max_concurrency = Keyword.get(nested_config, :max_concurrency, default_concurrency)
timeout = Keyword.get(nested_config, :timeout, default_timeout)
results =
libs_and_functions
|> Task.async_stream(
fn {library, python_module, functions} ->
{library, introspect_symbols(library, functions, python_module), python_module}
end,
max_concurrency: max_concurrency,
timeout: timeout
)
|> Enum.to_list()
libs_and_functions
|> Enum.zip(results)
|> Enum.map(fn
{{_library, _python_module, _functions}, {:ok, result}} ->
result
{{library, python_module, functions}, {:exit, reason}} ->
{library, {:error, batch_error(library, python_module, functions, reason)}, python_module}
{{library, python_module, functions}, {:error, reason}} ->
{library, {:error, batch_error(library, python_module, functions, reason)}, python_module}
end)
end
@doc """
Introspects a single attribute on a module to determine its type.
"""
@spec introspect_attribute(String.t() | atom(), String.t() | atom(), keyword()) ::
{:ok, map()} | {:error, term()}
def introspect_attribute(module_path, attr_name, opts \\ []) do
runner_opts = Keyword.merge(runner_opts(), opts)
case run_script(
[
script_path(),
"--module",
to_string(module_path),
"--attribute",
to_string(attr_name)
],
runner_opts
) do
{output, 0} -> parse_attribute_output(output)
{output, _status} -> handle_python_error(output, to_string(module_path))
end
end
defp library_python_name(%{python_name: python_name}) when is_binary(python_name),
do: python_name
defp library_python_name(%{name: name}) when is_atom(name), do: Atom.to_string(name)
defp library_python_name(name) when is_atom(name), do: Atom.to_string(name)
defp library_python_name(name) when is_binary(name), do: name
defp library_python_name(_), do: "unknown"
defp script_path do
Path.join(to_string(:code.priv_dir(:snakebridge)), "python/introspect.py")
end
defp library_label(%{name: name}) when is_atom(name), do: name
defp library_label(name) when is_atom(name), do: name
defp library_label(_), do: :unknown
defp runner_opts do
config = Application.get_env(:snakebridge, :introspector, [])
Keyword.take(config, [:timeout, :env, :cd])
end
defp parse_output(output) do
case Jason.decode(output) do
{:ok, results} when is_list(results) -> {:ok, results}
{:ok, %{"error" => error}} -> {:error, error}
{:error, _} -> {:error, {:json_parse, output}}
end
end
defp parse_attribute_output(output) do
case Jason.decode(output) do
{:ok, %{"error" => error}} -> {:error, error}
{:ok, result} when is_map(result) -> {:ok, result}
{:error, _} -> {:error, {:json_parse, output}}
end
end
defp group_symbols(infos) do
infos
|> Enum.reduce(%{"functions" => [], "classes" => [], "attributes" => []}, fn info, acc ->
case info["type"] || info[:type] do
"class" ->
Map.update!(acc, "classes", &[info | &1])
"attribute" ->
Map.update!(acc, "attributes", &[info | &1])
_ ->
Map.update!(acc, "functions", &[info | &1])
end
end)
|> Map.update!("functions", &Enum.reverse/1)
|> Map.update!("classes", &Enum.reverse/1)
|> Map.update!("attributes", &Enum.reverse/1)
end
defp run_script(args, opts) do
{timeout, opts} = Keyword.pop(opts, :timeout)
cmd_opts = [stderr_to_stdout: true]
env = build_env(opts)
cmd_opts = if env == [], do: cmd_opts, else: Keyword.put(cmd_opts, :env, env)
cmd_opts = maybe_put_opt(cmd_opts, :cd, opts)
run = fn -> System.cmd(python_executable(), args, cmd_opts) end
if is_integer(timeout) do
task = Task.async(run)
case Task.yield(task, timeout) || Task.shutdown(task, :brutal_kill) do
{:ok, result} -> result
nil -> {"Command timed out after #{timeout}ms", 124}
end
else
run.()
end
end
defp python_executable do
Application.get_env(:snakebridge, :python_executable) ||
resolve_snakepit_executable() ||
System.find_executable("python3") ||
"python3"
end
defp resolve_snakepit_executable do
if Code.ensure_loaded?(Snakepit.PythonRuntime) and
function_exported?(Snakepit.PythonRuntime, :resolve_executable, 0) do
case Snakepit.PythonRuntime.resolve_executable() do
{:ok, python, _meta} -> python
_ -> nil
end
else
nil
end
end
defp build_env(opts) do
runtime_env =
if Code.ensure_loaded?(Snakepit.PythonRuntime) and
function_exported?(Snakepit.PythonRuntime, :runtime_env, 0) do
Snakepit.PythonRuntime.runtime_env()
else
[]
end
extra_env =
if Code.ensure_loaded?(Snakepit.PythonRuntime) and
function_exported?(Snakepit.PythonRuntime, :config, 0) do
Snakepit.PythonRuntime.config() |> Map.get(:extra_env, %{}) |> Enum.to_list()
else
[]
end
user_env =
opts
|> Keyword.get(:env, %{})
|> Enum.to_list()
runtime_env ++ extra_env ++ user_env
end
defp maybe_put_opt(cmd_opts, key, opts) do
case Keyword.get(opts, key) do
nil -> cmd_opts
value -> Keyword.put(cmd_opts, key, value)
end
end
defp handle_python_error(output, package) do
{:error, IntrospectionError.from_python_output(output, package)}
end
defp batch_error(library, python_module, functions, reason) do
%{
type: :introspection_batch_failed,
library: library_label(library),
python_module: python_module,
functions: functions,
reason: reason
}
end
end
defmodule SnakeBridge.InvalidRefError do
@moduledoc """
Raised when a ref payload is malformed or invalid.
This occurs when the ref structure is missing required fields or has
an unrecognized format.
## Fields
- `:reason` - Why the ref is invalid (atom or string)
- `:message` - Human-readable error message
"""
defexception [:reason, :message]
@type t :: %__MODULE__{
reason: atom() | String.t() | nil,
message: String.t()
}
@impl Exception
def exception(opts) when is_list(opts) do
reason = Keyword.get(opts, :reason)
message = Keyword.get(opts, :message) || build_message(reason)
%__MODULE__{
reason: reason,
message: message
}
end
@impl Exception
def message(%__MODULE__{message: message}), do: message
defp build_message(reason) when is_atom(reason) do
case reason do
:missing_id -> "Invalid SnakeBridge reference: missing 'id' field"
:missing_type -> "Invalid SnakeBridge reference: missing '__type__' field"
:invalid_format -> "Invalid SnakeBridge reference: unrecognized payload format"
_ -> "Invalid SnakeBridge reference: #{reason}"
end
end
defp build_message(reason) when is_binary(reason) do
"Invalid SnakeBridge reference: #{reason}"
end
defp build_message(_) do
"Invalid SnakeBridge reference"
end
end
defmodule SnakeBridge.Ledger do
@moduledoc """
Wrapper for recording dynamic calls through Snakepit.
"""
@spec dynamic_call(atom() | String.t(), atom() | String.t(), list(), keyword()) ::
{:ok, term()} | {:error, term()}
def dynamic_call(library, function, args, opts \\ []) do
if function_exported?(Snakepit, :dynamic_call, 4) do
# credo:disable-for-next-line Credo.Check.Refactor.Apply
apply(Snakepit, :dynamic_call, [library, function, args, opts])
else
{:error, :snakepit_dynamic_call_unavailable}
end
end
end
defmodule SnakeBridge.Lock do
@moduledoc """
Manages snakebridge.lock with runtime identity, hardware info, and library versions.
The lock file captures:
- Hardware identity (accelerator, CUDA version, CPU features)
- Platform information (OS, architecture)
- Python environment (version, packages)
- Library configurations
## Hardware-Aware Lock Files
The lock file includes hardware information to detect compatibility issues:
%{
"environment" => %{
"hardware" => %{
"accelerator" => "cuda",
"cuda_version" => "12.1",
"gpu_count" => 2,
"cpu_features" => ["avx", "avx2"]
},
"platform" => %{
"os" => "linux",
"arch" => "x86_64"
}
}
}
Use `SnakeBridge.Lock.Verifier` to verify compatibility.
"""
alias SnakeBridge.Config
@spec load() :: map() | nil
def load do
case File.read(lock_path()) do
{:ok, content} -> Jason.decode!(content)
{:error, :enoent} -> nil
end
end
@spec update(SnakeBridge.Config.t()) :: :ok
def update(config) do
lock = build(config)
lock
|> Jason.encode!(pretty: true)
|> then(&File.write!(lock_path(), &1))
end
@spec build(SnakeBridge.Config.t()) :: map()
def build(config) do
runtime =
case python_runtime_module().runtime_identity() do
{:ok, identity} -> identity
{:error, _} -> %{version: "unknown", platform: "unknown", hash: "unknown"}
end
packages = get_package_metadata(config)
packages_hash = "sha256:" <> compute_packages_hash(packages)
hardware = build_hardware_section()
platform = build_platform_section()
%{
"version" => version(),
"environment" => %{
"snakebridge_version" => version(),
"generator_hash" => generator_hash(),
"python_version" => runtime.version,
"python_platform" => runtime.platform,
"python_runtime_hash" => runtime.hash,
"python_packages_hash" => packages_hash,
"elixir_version" => System.version(),
"otp_version" => System.otp_release(),
"hardware" => hardware,
"platform" => platform
},
"compatibility" => build_compatibility_section(hardware),
"libraries" => libraries_lock(config),
"python_packages" => packages
}
end
@doc """
Builds the hardware section for the lock file.
Returns a map with hardware identity including accelerator type,
CUDA version if available, GPU count, and CPU features.
"""
@spec build_hardware_section() :: map()
def build_hardware_section do
identity = hardware_module().identity()
caps = hardware_module().capabilities()
base = %{
"accelerator" => identity["accelerator"],
"gpu_count" => identity["gpu_count"],
"cpu_features" => identity["cpu_features"]
}
# Add CUDA-specific info if available
base =
if caps.cuda do
base
|> Map.put("cuda_version", caps.cuda_version)
|> Map.put("cudnn_version", caps.cudnn_version)
else
base
end
base
end
@doc """
Builds the platform section for the lock file.
"""
@spec build_platform_section() :: map()
def build_platform_section do
identity = hardware_module().identity()
platform = identity["platform"] || "unknown-unknown"
case String.split(platform, "-", parts: 2) do
[os, arch] ->
%{"os" => os, "arch" => arch}
[os] ->
%{"os" => os, "arch" => "unknown"}
_ ->
%{"os" => "unknown", "arch" => "unknown"}
end
end
@doc """
Builds the compatibility section with minimum requirements.
"""
@spec build_compatibility_section(map()) :: map()
def build_compatibility_section(hardware) do
%{
"cuda_min" => hardware["cuda_version"],
"compute_capability_min" => nil
}
end
@doc """
Deterministic hash from sorted package versions.
"""
@spec compute_packages_hash(map()) :: String.t()
def compute_packages_hash(packages) when is_map(packages) do
packages
|> Enum.sort_by(fn {name, _} -> name end)
|> Enum.map_join("\n", fn {name, info} ->
version = Map.get(info, "version") || Map.get(info, :version) || "unknown"
"#{name}==#{version}"
end)
|> then(&:crypto.hash(:sha256, &1))
|> Base.encode16(case: :lower)
end
@doc """
Gets package metadata for the lockfile.
"""
@spec get_package_metadata(Config.t()) :: map()
def get_package_metadata(config) do
requirements = SnakeBridge.PythonEnv.derive_requirements(config.libraries)
if requirements == [] do
%{}
else
case python_packages_module().lock_metadata(requirements, python_packages_opts([])) do
{:ok, metadata} when is_map(metadata) -> metadata
{:error, _} -> %{}
end
end
end
defp libraries_lock(config) do
config.libraries
|> Enum.map(fn library ->
{
library.python_name,
%{
"requested" => library.version,
"resolved" => library.version,
"hash" => nil
}
}
end)
|> Map.new()
end
defp lock_path do
"snakebridge.lock"
end
defp version do
Application.spec(:snakebridge, :vsn) |> to_string()
end
@generator_files [
"lib/snakebridge/generator.ex",
"lib/snakebridge/docs.ex",
"priv/python/snakebridge_types.py",
"priv/python/snakebridge_adapter.py"
]
@doc """
Computes the generator hash from generator and adapter source contents.
"""
@spec generator_hash() :: String.t()
def generator_hash do
content = Enum.map_join(@generator_files, "\n", &read_generator_file/1)
:crypto.hash(:sha256, content)
|> Base.encode16(case: :lower)
end
defp read_generator_file(relative_path) do
candidates =
[
Application.app_dir(:snakebridge, relative_path),
Path.join(File.cwd!(), relative_path)
]
|> Enum.uniq()
Enum.find_value(candidates, "", fn path ->
case File.read(path) do
{:ok, content} -> content
{:error, _} -> nil
end
end) || ""
end
@doc """
Checks if the lock was generated with the current generator version.
"""
@spec verify_generator_unchanged?(map()) :: boolean()
def verify_generator_unchanged?(lock) do
lock_hash = get_in(lock, ["environment", "generator_hash"])
current_hash = generator_hash()
lock_hash == current_hash
end
defp python_packages_module do
Application.get_env(:snakebridge, :python_packages, Snakepit.PythonPackages)
end
defp python_packages_opts(opts) do
if python_packages_module() == Snakepit.PythonPackages do
Keyword.put_new(opts, :runner, SnakeBridge.PythonPackagesRunner)
else
opts
end
end
defp python_runtime_module do
Application.get_env(:snakebridge, :python_runtime, Snakepit.PythonRuntime)
end
defp hardware_module do
Application.get_env(:snakebridge, :hardware_module, Snakepit.Hardware)
end
end
defmodule SnakeBridge.Manifest do
@moduledoc """
Manifest storage for generated symbols.
"""
@spec load(SnakeBridge.Config.t()) :: map()
def load(config) do
path = manifest_path(config)
case File.read(path) do
{:ok, content} ->
content
|> Jason.decode!()
|> normalize_manifest()
{:error, :enoent} ->
%{"version" => version(), "symbols" => %{}, "classes" => %{}}
end
end
@spec save(SnakeBridge.Config.t(), map()) :: :ok
def save(config, manifest) do
path = manifest_path(config)
File.mkdir_p!(Path.dirname(path))
manifest
|> sort_manifest()
|> Jason.encode!(pretty: true)
|> then(&File.write!(path, &1))
end
@spec missing(map(), list({module(), atom(), non_neg_integer()})) ::
list({module(), atom(), non_neg_integer()})
def missing(manifest, detected) do
classes = Map.get(manifest, "classes", %{})
detected
|> Enum.reject(fn {mod, func, arity} ->
module_key = module_to_string(mod)
case Map.get(classes, module_key) do
nil -> call_supported?(manifest, mod, func, arity)
class_info -> class_call_supported?(class_info, func, arity)
end
end)
end
@spec call_supported?(map(), module(), atom(), non_neg_integer()) :: boolean()
def call_supported?(manifest, module, function, call_site_arity) do
prefix = "#{module_to_string(module)}.#{function}/"
manifest
|> Map.get("symbols", %{})
|> Enum.any?(fn {key, info} ->
String.starts_with?(key, prefix) and
symbol_arity_matches?(key, info, call_site_arity)
end)
end
defp symbol_arity_matches?(key, info, call_site_arity) do
arity_from_key = symbol_arity_from_key(key)
min_arity = info["minimum_arity"] || info["required_arity"] || arity_from_key || 0
max_arity = info["maximum_arity"] || arity_from_key
has_var_positional = info["has_var_positional"] == true
arity_in_range?(call_site_arity, min_arity, max_arity, has_var_positional)
end
defp arity_in_range?(call_site_arity, min_arity, max_arity, has_var_positional) do
cond do
max_arity in [:unbounded, "unbounded"] or has_var_positional ->
call_site_arity >= min_arity
is_integer(max_arity) ->
call_site_arity >= min_arity and call_site_arity <= max_arity
true ->
call_site_arity == min_arity
end
end
defp class_call_supported?(class_info, function, call_site_arity) do
function_name = to_string(function)
methods = method_field(class_info, "methods") || []
attrs = method_field(class_info, "attributes") || []
if methods == [] and attrs == [] do
true
else
method_supported? =
Enum.any?(methods, fn method ->
method_name(method) == function_name and
method_arity_supported?(method, call_site_arity)
end)
attr_supported? =
Enum.any?(attrs, fn attr ->
to_string(attr) == function_name and call_site_arity == 1
end)
method_supported? or attr_supported?
end
end
defp method_name(method) do
method_field(method, "elixir_name") ||
case method_field(method, "name") do
"__init__" -> "new"
name when is_binary(name) -> name
_ -> ""
end
end
defp method_arity_supported?(method, call_site_arity) do
{min_arity, max_arity, has_var_positional} = method_arity_info(method)
arity_in_range?(call_site_arity, min_arity, max_arity, has_var_positional)
end
defp method_arity_info(method) do
min_arity = method_field(method, "minimum_arity")
max_arity = method_field(method, "maximum_arity")
required_arity = method_field(method, "required_arity")
has_var_positional = method_field(method, "has_var_positional") == true
if has_explicit_arity_info?(min_arity, max_arity, has_var_positional) do
{min_arity || required_arity || 0, max_arity, has_var_positional}
else
compute_arity_from_params(method)
end
end
defp has_explicit_arity_info?(min_arity, max_arity, has_var_positional) do
is_integer(min_arity) or is_integer(max_arity) or has_var_positional
end
defp compute_arity_from_params(method) do
params = method_field(method, "parameters") || []
signature_available = method_field(method, "signature_available") != false
raw_name = method_field(method, "name") || ""
{min_base, max_base, var_positional?} =
compute_method_arity(params, signature_available)
ref_offset = if raw_name == "__init__", do: 0, else: 1
apply_ref_offset(min_base, max_base, ref_offset, var_positional?)
end
defp apply_ref_offset(min_base, max_base, ref_offset, var_positional?) do
min_arity = min_base + ref_offset
max_arity =
case max_base do
:unbounded -> :unbounded
value when is_integer(value) -> value + ref_offset
_ -> min_arity
end
{min_arity, max_arity, var_positional?}
end
defp compute_method_arity(params, signature_available) do
positional_params =
Enum.filter(params, fn param ->
param_kind(param) in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD"]
end)
required_positional =
positional_params
|> Enum.reject(¶m_default?/1)
|> length()
optional_positional =
positional_params
|> Enum.filter(¶m_default?/1)
|> length()
has_var_positional = Enum.any?(params, &varargs?/1)
variadic_fallback = params == [] and signature_available == false
max_arity =
cond do
variadic_fallback -> variadic_max_arity() + 1
has_var_positional -> :unbounded
optional_positional > 0 -> required_positional + 2
true -> required_positional + 1
end
{required_positional, max_arity, has_var_positional}
end
defp variadic_max_arity do
Application.get_env(:snakebridge, :variadic_max_arity, 8)
end
defp varargs?(param), do: param_kind(param) == "VAR_POSITIONAL"
defp param_kind(%{"kind" => kind}) when is_binary(kind), do: String.upcase(kind)
defp param_kind(%{kind: kind}) when is_binary(kind), do: String.upcase(kind)
defp param_kind(%{kind: kind}), do: kind
defp param_kind(_), do: nil
defp param_default?(%{"default" => _}), do: true
defp param_default?(%{default: _}), do: true
defp param_default?(_), do: false
defp method_field(map, key) when is_binary(key) do
Map.get(map, key) || Map.get(map, String.to_atom(key))
end
defp symbol_arity_from_key(key) when is_binary(key) do
case String.split(key, "/") do
[_prefix, arity] ->
case Integer.parse(arity) do
{value, ""} -> value
_ -> nil
end
_ ->
nil
end
end
@spec put_symbols(map(), list({String.t(), map()})) :: map()
def put_symbols(manifest, entries) do
symbols =
manifest
|> Map.get("symbols", %{})
|> Map.merge(Map.new(entries))
Map.put(manifest, "symbols", symbols)
end
@spec put_classes(map(), list({String.t(), map()})) :: map()
def put_classes(manifest, entries) do
classes =
manifest
|> Map.get("classes", %{})
|> Map.merge(Map.new(entries))
Map.put(manifest, "classes", classes)
end
@spec symbol_key({module(), atom(), non_neg_integer()}) :: String.t()
def symbol_key({module, function, arity}) do
mod = module |> Module.split() |> Enum.join(".")
"#{mod}.#{function}/#{arity}"
end
@spec class_key(module()) :: String.t()
def class_key(module) when is_atom(module) do
Module.split(module) |> Enum.join(".")
end
defp module_to_string(module) when is_atom(module) do
Module.split(module) |> Enum.join(".")
end
defp manifest_path(config) do
Path.join(config.metadata_dir, "manifest.json")
end
defp version do
Application.spec(:snakebridge, :vsn) |> to_string()
end
defp normalize_manifest(manifest) do
symbols = Map.get(manifest, "symbols", %{})
normalized_symbols =
Enum.reduce(symbols, %{}, fn {key, value}, acc ->
normalized = normalize_symbol_key(key)
if normalized == key do
Map.put(acc, normalized, value)
else
Map.put_new(acc, normalized, value)
end
end)
Map.put(manifest, "symbols", normalized_symbols)
end
defp normalize_symbol_key(key) when is_binary(key) do
case String.split(key, ".") do
["Elixir" | rest] ->
case Enum.split(rest, -1) do
{module_parts, [fun_part]} when module_parts != [] ->
Enum.join(module_parts, ".") <> "." <> fun_part
_ ->
key
end
_ ->
key
end
end
defp normalize_symbol_key(key), do: key
defp sort_manifest(manifest) do
manifest
|> update_in(["symbols"], fn symbols ->
symbols
|> Enum.sort_by(fn {key, _} -> key end)
|> Map.new()
end)
|> update_in(["classes"], fn classes ->
classes
|> Enum.sort_by(fn {key, _} -> key end)
|> Map.new()
end)
end
end
defmodule SnakeBridge.ModuleResolver do
@moduledoc """
Resolves ambiguous module paths to class attributes or submodules.
"""
alias SnakeBridge.Introspector
@type resolution ::
{:class, String.t(), String.t()}
| {:submodule, String.t()}
| {:error, term()}
@doc """
Determines if an Elixir module maps to a Python class attribute or submodule.
Returns:
- `{:class, class_name, parent_module}` when the last path segment is a class.
- `{:submodule, module_path}` when the path resolves to a submodule.
- `{:error, reason}` when introspection fails.
"""
@spec resolve_class_or_submodule(map(), module()) :: resolution()
def resolve_class_or_submodule(library, elixir_module) do
module_parts = Module.split(elixir_module)
library_parts = Module.split(library.module_name)
extra_parts = Enum.drop(module_parts, length(library_parts))
case extra_parts do
[] ->
{:submodule, library.python_name}
_ ->
{parent_parts, [candidate]} = Enum.split(extra_parts, -1)
parent_module = build_parent_module(library.python_name, parent_parts)
case class_or_module(parent_module, candidate) do
{:class, class_name} ->
{:class, class_name, parent_module}
{:module, module_name} ->
{:submodule, join_module(parent_module, module_name)}
:unknown ->
{:submodule, fallback_submodule(library.python_name, extra_parts)}
{:error, _} = error ->
error
end
end
end
defp class_or_module(parent_module, candidate) do
case introspect_attribute_type(parent_module, candidate) do
{:ok, :class} ->
{:class, candidate}
{:ok, :module} ->
{:module, candidate}
{:ok, :other} ->
maybe_downcase(parent_module, candidate)
{:error, _} = error ->
error
end
end
defp maybe_downcase(parent_module, candidate) do
downcased = String.downcase(candidate)
if downcased == candidate do
:unknown
else
case introspect_attribute_type(parent_module, downcased) do
{:ok, :class} -> {:class, downcased}
{:ok, :module} -> {:module, downcased}
{:ok, :other} -> :unknown
{:error, _} = error -> error
end
end
end
defp introspect_attribute_type(module_path, attr_name) do
case Introspector.introspect_attribute(module_path, attr_name) do
{:ok, %{"exists" => false}} ->
{:ok, :other}
{:ok, %{"is_class" => true}} ->
{:ok, :class}
{:ok, %{"is_module" => true}} ->
{:ok, :module}
{:ok, _} ->
{:ok, :other}
{:error, _} = error ->
error
end
end
defp build_parent_module(base, []), do: base
defp build_parent_module(base, parts) do
base <> "." <> Enum.map_join(parts, ".", &Macro.underscore/1)
end
defp fallback_submodule(base, parts) do
[base | Enum.map(parts, &Macro.underscore/1)]
|> Enum.join(".")
end
defp join_module(parent, child), do: parent <> "." <> child
end
defmodule SnakeBridge.PythonEnv do
@moduledoc """
Compile-time orchestrator for Python environment provisioning.
"""
alias SnakeBridge.{Config, EnvironmentError}
@type requirement :: String.t()
@doc """
Ensures the Python environment is ready for introspection.
In dev with auto_install enabled, installs missing packages.
In strict mode, verifies the environment without installing.
"""
@spec ensure!(Config.t()) :: :ok | no_return()
def ensure!(config) do
cond do
strict_mode?(config) ->
verify_environment!(config)
auto_install_enabled?(config) ->
do_ensure!(config)
true ->
verify_environment!(config)
end
end
@doc """
Converts library config to PEP-440 requirement strings.
Skips stdlib libraries and applies pypi_package and extras overrides.
"""
@spec derive_requirements([Config.Library.t()]) :: [requirement()]
def derive_requirements(libraries) when is_list(libraries) do
libraries
|> Enum.reject(&stdlib_library?/1)
|> Enum.map(&library_to_requirement/1)
end
@doc """
Checks packages are installed without installing.
"""
@spec verify_environment!(Config.t()) :: :ok | no_return()
def verify_environment!(config) do
requirements = derive_requirements(config.libraries)
case python_packages_module().check_installed(requirements, python_packages_opts([])) do
{:ok, :all_installed} ->
:ok
{:ok, {:missing, missing}} ->
raise EnvironmentError,
message: "Missing Python packages: #{inspect(missing)}",
missing_packages: missing,
suggestion: "Run: mix snakebridge.setup"
end
end
defp do_ensure!(config) do
ensure_python_runtime!()
ensure_snakepit_adapter!()
ensure_snakepit_requirements!(config)
requirements = derive_requirements(config.libraries)
if requirements != [] do
python_packages_module().ensure!(
{:list, requirements},
python_packages_opts(quiet: !config.verbose)
)
end
:ok
end
defp ensure_snakepit_requirements!(config) do
if python_packages_module() == Snakepit.PythonPackages do
case snakepit_requirements_path() do
nil ->
:ok
path ->
python_packages_module().ensure!(
{:file, path},
python_packages_opts(quiet: !config.verbose)
)
end
else
:ok
end
end
defp ensure_python_runtime! do
python_config = Application.get_env(:snakepit, :python, [])
if Keyword.get(python_config, :managed, false) do
python_runtime_module().install_managed(SnakeBridge.PythonRuntimeRunner, [])
end
:ok
end
defp auto_install_enabled?(config) do
case auto_install_setting(config) do
:never -> false
:always -> true
:dev -> Mix.env() == :dev
end
end
defp auto_install_setting(config) do
case System.get_env("SNAKEBRIDGE_AUTO_INSTALL") do
"never" -> :never
"always" -> :always
"dev" -> :dev
nil -> config.auto_install || :dev
_ -> config.auto_install || :dev
end
end
defp strict_mode?(config) do
System.get_env("SNAKEBRIDGE_STRICT") == "1" || config.strict == true
end
defp stdlib_library?(%Config.Library{version: :stdlib}), do: true
defp stdlib_library?(_), do: false
defp library_to_requirement(library) do
package = library.pypi_package || library.python_name || Atom.to_string(library.name)
extras = List.wrap(library.extras || [])
version = translate_version(library.version)
base =
if extras == [] do
package
else
package <> "[" <> Enum.join(extras, ",") <> "]"
end
if version do
base <> version
else
base
end
end
defp translate_version(nil), do: nil
defp translate_version(:stdlib), do: nil
defp translate_version(v) when is_binary(v) do
v = String.trim(v)
if String.starts_with?(v, ["~=", ">=", "<=", "==", "!="]) do
v
else
case Regex.run(~r/^~>\s*(.+)$/, v) do
[_, ver] -> "~=#{ver}"
nil -> "==#{v}"
end
end
end
defp python_packages_module do
Application.get_env(:snakebridge, :python_packages, Snakepit.PythonPackages)
end
defp python_packages_opts(opts) do
if python_packages_module() == Snakepit.PythonPackages do
Keyword.put_new(opts, :runner, SnakeBridge.PythonPackagesRunner)
else
opts
end
end
defp python_runtime_module do
Application.get_env(:snakebridge, :python_runtime, Snakepit.PythonRuntime)
end
defp snakepit_requirements_path do
case :code.priv_dir(:snakepit) do
{:error, _} ->
nil
priv_dir ->
path = Path.join([to_string(priv_dir), "python", "requirements.txt"])
if File.exists?(path), do: path, else: nil
end
end
defp ensure_snakepit_adapter! do
if is_list(Application.get_env(:snakepit, :pools)) do
:ok
else
adapter_module = Application.get_env(:snakepit, :adapter_module)
if is_nil(adapter_module) do
Application.put_env(:snakepit, :adapter_module, Snakepit.Adapters.GRPCPython)
end
pool_config =
:snakepit
|> Application.get_env(:pool_config, %{})
|> normalize_config_input()
adapter_args = Map.get(pool_config, :adapter_args, [])
if adapter_args_missing?(adapter_args) do
updated = Map.put(pool_config, :adapter_args, default_adapter_args())
Application.put_env(:snakepit, :pool_config, updated)
end
:ok
end
end
defp adapter_args_missing?(adapter_args) when is_list(adapter_args) do
not Enum.any?(adapter_args, fn arg ->
is_binary(arg) and (arg == "--adapter" or String.starts_with?(arg, "--adapter="))
end)
end
defp adapter_args_missing?(_), do: true
defp default_adapter_args do
["--adapter", "snakebridge_adapter.SnakeBridgeAdapter"]
end
defp normalize_config_input(nil), do: %{}
defp normalize_config_input(%{} = map), do: map
defp normalize_config_input(list) when is_list(list), do: Map.new(list)
defp normalize_config_input(_), do: %{}
end
defmodule SnakeBridge.PythonPackagesRunner do
@moduledoc false
@behaviour Snakepit.PythonPackages.Runner
@impl true
@doc false
def cmd(command, args, opts) do
{timeout, opts} = Keyword.pop(opts, :timeout)
run = fn -> System.cmd(command, args, opts) end
if is_integer(timeout) do
task = Task.async(run)
case Task.yield(task, timeout) || Task.shutdown(task, :brutal_kill) do
{:ok, result} -> result
nil -> {"Command timed out after #{timeout}ms", 124}
end
else
run.()
end
rescue
error in ErlangError ->
{Exception.message(error), 127}
end
end
defmodule SnakeBridge.PythonRunner do
@moduledoc """
Behaviour for executing Python scripts in the Snakepit-configured runtime.
"""
@type script :: String.t()
@type args :: [String.t()]
@type opts :: keyword()
@callback run(script(), args(), opts()) :: {:ok, String.t()} | {:error, term()}
end
defmodule SnakeBridge.PythonRuntimeRunner do
@moduledoc false
@behaviour Snakepit.Bootstrap.Runner
@impl true
def mix(_task, _args), do: :ok
@impl true
def cmd(command, args, opts) do
opts = Keyword.put_new(opts, :stderr_to_stdout, true)
{actual_command, actual_args} =
if String.ends_with?(command, ".sh") do
{"bash", [command | args]}
else
{command, args}
end
case System.cmd(actual_command, actual_args, opts) do
{output, 0} ->
write_output(output)
:ok
{output, status} ->
write_output(output)
{:error, {:command_failed, command, status}}
end
rescue
error in ErlangError ->
{:error, {:command_failed, command, Exception.message(error)}}
end
defp write_output(output) do
if String.trim(output) != "" do
IO.write(output)
end
end
end
defmodule SnakeBridge.RefNotFoundError do
@moduledoc """
Raised when a Python object reference cannot be found in the registry.
This typically occurs when:
- The ref was already released via `release_ref/1`
- The session was released via `release_session/1`
- The ref expired due to TTL
- The ref was evicted due to registry size limits
## Fields
- `:ref_id` - The ref ID that was not found
- `:session_id` - The session ID the ref was looked up in
- `:message` - Human-readable error message
"""
defexception [:ref_id, :session_id, :message]
@type t :: %__MODULE__{
ref_id: String.t() | nil,
session_id: String.t() | nil,
message: String.t()
}
@impl Exception
def exception(opts) when is_list(opts) do
ref_id = Keyword.get(opts, :ref_id)
session_id = Keyword.get(opts, :session_id)
message = Keyword.get(opts, :message) || build_message(ref_id, session_id)
%__MODULE__{
ref_id: ref_id,
session_id: session_id,
message: message
}
end
@impl Exception
def message(%__MODULE__{message: message}), do: message
defp build_message(ref_id, session_id) do
base = "SnakeBridge reference '#{ref_id || "unknown"}' not found"
if session_id do
base <> " in session '#{session_id}'. The ref may have been released, expired, or evicted."
else
base <> ". The ref may have been released, expired, or evicted."
end
end
end
defmodule SnakeBridge.Ref do
@moduledoc """
Structured reference to a Python object managed by SnakeBridge.
This struct defines the cross-language wire shape for Python object references.
"""
@schema_version 1
@typedoc """
Structured reference to a Python object.
"""
@type t :: %__MODULE__{
id: String.t(),
session_id: String.t(),
python_module: String.t() | nil,
library: String.t() | nil,
type_name: String.t() | nil,
schema: pos_integer()
}
defstruct [
:id,
:session_id,
:python_module,
:library,
:type_name,
schema: @schema_version
]
@spec schema_version() :: pos_integer()
def schema_version, do: @schema_version
@doc """
Creates a Ref from a wire format map.
"""
@spec from_wire_format(map() | t()) :: t()
def from_wire_format(%__MODULE__{} = ref), do: ref
def from_wire_format(map) when is_map(map) do
%__MODULE__{
id: get_wire_field(map, ["id", "ref_id"]),
session_id: get_wire_field(map, ["session_id"]),
python_module: get_wire_field(map, ["python_module"]),
library: get_wire_field(map, ["library"]),
type_name: get_wire_field(map, ["type_name", "__type_name__"]),
schema: get_wire_field(map, ["__schema__", "schema"]) || @schema_version
}
end
defp get_wire_field(map, keys) do
Enum.find_value(keys, fn key ->
Map.get(map, key) || Map.get(map, String.to_atom(key))
end)
end
@doc """
Converts a Ref to wire format for Python calls.
"""
@spec to_wire_format(t() | map()) :: map()
def to_wire_format(%__MODULE__{} = ref) do
%{}
|> Map.put("__type__", "ref")
|> Map.put("__schema__", ref.schema || @schema_version)
|> maybe_put("id", ref.id)
|> maybe_put("session_id", ref.session_id)
|> maybe_put("python_module", ref.python_module)
|> maybe_put("library", ref.library)
end
def to_wire_format(%{"__type__" => "ref"} = ref) do
ref
|> from_wire_format()
|> to_wire_format()
end
def to_wire_format(%{__type__: "ref"} = ref) do
ref
|> from_wire_format()
|> to_wire_format()
end
@doc """
Checks if a value is a valid ref.
"""
@spec ref?(term()) :: boolean()
def ref?(%__MODULE__{id: id, session_id: session_id})
when is_binary(id) and is_binary(session_id),
do: true
def ref?(%{"__type__" => "ref"} = ref) do
ref_id =
Map.get(ref, "id") || Map.get(ref, :id) || Map.get(ref, "ref_id") || Map.get(ref, :ref_id)
session_id = Map.get(ref, "session_id") || Map.get(ref, :session_id)
is_binary(ref_id) and is_binary(session_id)
end
def ref?(%{__type__: "ref"} = ref), do: ref?(Map.new(ref, fn {k, v} -> {to_string(k), v} end))
def ref?(_), do: false
defp maybe_put(map, _key, nil), do: map
defp maybe_put(map, key, value), do: Map.put(map, key, value)
end
defimpl Inspect, for: SnakeBridge.Ref do
import Inspect.Algebra
alias SnakeBridge.Runtime
alias SnakeBridge.Ref
def inspect(%Ref{} = ref, _opts) do
case python_repr(ref) do
{:ok, repr} when is_binary(repr) ->
concat(["#Python<", repr, ">"])
_ ->
concat(["#SnakeBridge.Ref<", to_string(ref.id || "unknown"), ">"])
end
end
defp python_repr(ref) do
case safe_call(ref, :__repr__) do
{:ok, repr} when is_binary(repr) ->
{:ok, repr}
_ ->
python_str(ref)
end
end
defp python_str(ref) do
case safe_call(ref, :__str__) do
{:ok, str} when is_binary(str) -> {:ok, str}
_ -> {:error, :unavailable}
end
end
defp safe_call(ref, method) do
Runtime.call_method(ref, method, [])
rescue
exception -> {:error, exception}
end
end
defimpl String.Chars, for: SnakeBridge.Ref do
alias SnakeBridge.Runtime
alias SnakeBridge.Ref
def to_string(%Ref{} = ref) do
case safe_call(ref, :__str__) do
{:ok, str} when is_binary(str) -> str
_ -> "#SnakeBridge.Ref<#{ref.id || "unknown"}>"
end
end
defp safe_call(ref, method) do
Runtime.call_method(ref, method, [])
rescue
exception -> {:error, exception}
end
end
defimpl Enumerable, for: SnakeBridge.Ref do
alias SnakeBridge.Runtime
alias SnakeBridge.StreamRef
def count(ref) do
case safe_call(ref, :__len__, []) do
{:ok, len} when is_integer(len) -> {:ok, len}
_ -> {:error, __MODULE__}
end
end
def member?(ref, value) do
case safe_call(ref, :__contains__, [value]) do
{:ok, result} when is_boolean(result) -> {:ok, result}
_ -> {:error, __MODULE__}
end
end
def slice(_ref), do: {:error, __MODULE__}
def reduce(ref, acc, fun) do
case safe_call(ref, :__iter__, []) do
{:ok, %StreamRef{} = stream_ref} ->
Enumerable.reduce(stream_ref, acc, fun)
{:ok, iterator_ref} ->
do_reduce(iterator_ref, acc, fun)
{:error, _} ->
do_reduce_by_index(ref, 0, acc, fun)
end
end
defp do_reduce(_iterator, {:halt, acc}, _fun), do: {:halted, acc}
defp do_reduce(iterator, {:suspend, acc}, fun) do
{:suspended, acc, &do_reduce(iterator, &1, fun)}
end
defp do_reduce(iterator, {:cont, acc}, fun) do
case safe_call(iterator, :__next__, []) do
{:ok, value} ->
do_reduce(iterator, fun.(value, acc), fun)
{:error, reason} ->
if stop_iteration?(reason) do
{:done, acc}
else
{:halted, {:error, reason}}
end
end
end
defp do_reduce_by_index(_ref, _index, {:halt, acc}, _fun), do: {:halted, acc}
defp do_reduce_by_index(ref, index, {:suspend, acc}, fun) do
{:suspended, acc, &do_reduce_by_index(ref, index, &1, fun)}
end
defp do_reduce_by_index(ref, index, {:cont, acc}, fun) do
case safe_call(ref, :__getitem__, [index]) do
{:ok, value} ->
do_reduce_by_index(ref, index + 1, fun.(value, acc), fun)
{:error, reason} ->
if index_error?(reason) or stop_iteration?(reason) do
{:done, acc}
else
{:halted, {:error, reason}}
end
end
end
defp safe_call(ref, method, args) do
Runtime.call_method(ref, method, args)
rescue
exception -> {:error, exception}
end
defp stop_iteration?(reason), do: error_type(reason) == "StopIteration"
defp index_error?(reason), do: error_type(reason) == "IndexError"
defp error_type(%{python_class: class}) when is_binary(class) do
class
|> String.split(".")
|> List.last()
end
# Check for python_type field (these are atom-keyed struct fields)
defp error_type(%{python_type: type}) when is_binary(type), do: type
defp error_type(%{python_type: type}) when is_atom(type), do: Atom.to_string(type)
# Check for error_type field
defp error_type(%{error_type: type}) when is_binary(type), do: type
defp error_type(%{error_type: type}) when is_atom(type), do: Atom.to_string(type)
# Check for type field
defp error_type(%{type: type}) when is_binary(type), do: type
defp error_type(%{type: type}) when is_atom(type), do: Atom.to_string(type)
# Handle exception structs by extracting the module name
defp error_type(%{__struct__: struct}) do
struct
|> Module.split()
|> List.last()
end
end
defmodule SnakeBridge.Registry do
@moduledoc """
Registry system for tracking generated SnakeBridge adapters.
The registry maintains a record of all generated Python library adapters,
allowing agents and tools to introspect what libraries are available without
parsing code.
## Registry Format
The registry stores library information including:
- Python module name and version
- Generated Elixir module name
- Generation timestamp
- File locations and structure
- Statistics (function count, class count, etc.)
## Usage
# Register a new library
SnakeBridge.Registry.register("numpy", %{
python_module: "numpy",
python_version: "1.26.0",
elixir_module: "Numpy",
generated_at: ~U[2024-12-24 14:00:00Z],
path: "lib/snakebridge/adapters/numpy/",
files: ["numpy.ex", "linalg.ex", "_meta.ex"],
stats: %{functions: 165, classes: 2, submodules: 4}
})
# Check if a library is generated
SnakeBridge.Registry.generated?("numpy")
# => true
# Get library information
SnakeBridge.Registry.get("numpy")
# => %{python_module: "numpy", ...}
# List all generated libraries
SnakeBridge.Registry.list_libraries()
# => ["json", "numpy", "sympy"]
## Persistence
The registry is automatically persisted to a JSON file at:
`priv/snakebridge/registry.json`
Use `save/0` to persist changes and `load/0` to restore from disk.
"""
use Agent
require Logger
@type library_name :: String.t()
@type registry_entry :: %{
python_module: String.t(),
python_version: String.t(),
elixir_module: String.t(),
generated_at: DateTime.t(),
path: String.t(),
files: [String.t()],
stats: %{
functions: non_neg_integer(),
classes: non_neg_integer(),
submodules: non_neg_integer()
}
}
@type registry_state :: %{
optional(library_name()) => registry_entry()
}
# Registry version for compatibility tracking
@registry_version "2.1"
# Required entry fields for validation
@required_fields [
:python_module,
:python_version,
:elixir_module,
:generated_at,
:path,
:files,
:stats
]
@required_stat_fields [:functions, :classes, :submodules]
## Client API
@doc """
Starts the registry agent.
This is typically called by the application supervisor.
"""
@spec start_link(keyword()) :: Agent.on_start()
def start_link(opts \\ []) do
Agent.start_link(fn -> %{} end, Keyword.merge([name: __MODULE__], opts))
end
@doc """
Returns a list of all registered library names, sorted alphabetically.
## Examples
iex> SnakeBridge.Registry.register("numpy", entry)
:ok
iex> SnakeBridge.Registry.list_libraries()
["numpy"]
"""
@spec list_libraries() :: [library_name()]
def list_libraries do
registry_get(fn state ->
state
|> Map.keys()
|> Enum.sort()
end)
end
@doc """
Gets information about a registered library.
Returns `nil` if the library is not registered.
## Examples
iex> SnakeBridge.Registry.get("numpy")
%{python_module: "numpy", python_version: "1.26.0", ...}
iex> SnakeBridge.Registry.get("nonexistent")
nil
"""
@spec get(library_name()) :: registry_entry() | nil
def get(library_name) do
registry_get(fn state -> Map.get(state, library_name) end)
end
@doc """
Checks if a library is registered.
## Examples
iex> SnakeBridge.Registry.generated?("numpy")
true
iex> SnakeBridge.Registry.generated?("nonexistent")
false
"""
@spec generated?(library_name()) :: boolean()
def generated?(library_name) do
registry_get(fn state -> Map.has_key?(state, library_name) end)
end
@doc """
Registers a library in the registry.
Updates the entry if the library is already registered.
## Parameters
- `library_name` - The library identifier (e.g., "numpy")
- `entry` - A map containing library information (see module documentation)
## Returns
- `:ok` on success
- `{:error, reason}` if the entry is invalid
## Examples
iex> entry = %{
...> python_module: "numpy",
...> python_version: "1.26.0",
...> elixir_module: "Numpy",
...> generated_at: ~U[2024-12-24 14:00:00Z],
...> path: "lib/snakebridge/adapters/numpy/",
...> files: ["numpy.ex"],
...> stats: %{functions: 10, classes: 0, submodules: 1}
...> }
iex> SnakeBridge.Registry.register("numpy", entry)
:ok
"""
@spec register(library_name(), map()) :: :ok | {:error, String.t()}
def register(library_name, entry) when is_binary(library_name) and is_map(entry) do
case validate_entry(entry) do
:ok ->
registry_update(fn state ->
Map.put(state, library_name, normalize_entry(entry))
end)
:ok
{:error, _reason} = error ->
error
end
end
@doc """
Removes a library from the registry.
Returns `:ok` even if the library was not registered.
## Examples
iex> SnakeBridge.Registry.unregister("numpy")
:ok
"""
@spec unregister(library_name()) :: :ok
def unregister(library_name) do
registry_update(fn state ->
Map.delete(state, library_name)
end)
:ok
end
@doc """
Clears all entries from the registry.
## Examples
iex> SnakeBridge.Registry.clear()
:ok
"""
@spec clear() :: :ok
def clear do
registry_update(fn _state -> %{} end)
:ok
end
@doc """
Saves the registry to the JSON file.
Creates the parent directory if it doesn't exist.
## Returns
- `:ok` on success
- `{:error, reason}` if saving fails
## Examples
iex> SnakeBridge.Registry.save()
:ok
"""
@spec save() :: :ok | {:error, term()}
def save do
registry_path = get_registry_path()
with :ok <- ensure_registry_dir(registry_path),
{:ok, data} <- build_registry_data(),
{:ok, json} <- Jason.encode(data, pretty: true),
:ok <- File.write(registry_path, json) do
:ok
else
{:error, reason} = error ->
Logger.error("Failed to save registry to #{registry_path}: #{inspect(reason)}")
error
end
end
@doc """
Loads the registry from the JSON file.
If the file doesn't exist, initializes an empty registry.
## Returns
- `:ok` on success
- `{:error, reason}` if loading fails
## Examples
iex> SnakeBridge.Registry.load()
:ok
"""
@spec load() :: :ok | {:error, term()}
def load do
registry_path = get_registry_path()
case File.read(registry_path) do
{:ok, content} ->
with {:ok, data} <- Jason.decode(content),
{:ok, libraries} <- parse_registry_data(data) do
registry_update(fn _state -> libraries end)
:ok
else
{:error, reason} = error ->
Logger.error("Failed to parse registry from #{registry_path}: #{inspect(reason)}")
error
end
{:error, :enoent} ->
# File doesn't exist yet - start with empty registry
Logger.debug("Registry file not found at #{registry_path}, starting with empty registry")
:ok
{:error, reason} = error ->
Logger.error("Failed to read registry from #{registry_path}: #{inspect(reason)}")
error
end
end
## Private Functions
# Ensures the registry agent is started
defp ensure_started do
case Process.whereis(__MODULE__) do
nil ->
start_registry()
pid ->
if Process.alive?(pid) do
:ok
else
start_registry()
end
end
end
defp start_registry do
case Agent.start(fn -> %{} end, name: __MODULE__) do
{:ok, _pid} -> :ok
{:error, {:already_started, _pid}} -> :ok
{:error, reason} -> raise "Failed to start registry: #{inspect(reason)}"
end
end
# Gets the registry file path from config or default
defp get_registry_path do
Application.get_env(:snakebridge, :registry_path) ||
Path.join([File.cwd!(), "priv", "snakebridge", "registry.json"])
end
# Ensures the registry directory exists
defp ensure_registry_dir(registry_path) do
registry_path
|> Path.dirname()
|> File.mkdir_p()
end
# Validates a registry entry has all required fields
defp validate_entry(entry) do
# Check required top-level fields
missing_fields =
@required_fields
|> Enum.reject(fn field -> Map.has_key?(entry, field) end)
cond do
length(missing_fields) > 0 ->
{:error, "Missing required fields: #{inspect(missing_fields)}"}
not is_map(entry.stats) ->
{:error, "stats must be a map"}
true ->
validate_stats(entry.stats)
end
end
# Validates the stats sub-map
defp validate_stats(stats) do
missing_stat_fields =
@required_stat_fields
|> Enum.reject(fn field -> Map.has_key?(stats, field) end)
if length(missing_stat_fields) > 0 do
{:error, "Missing required stat fields: #{inspect(missing_stat_fields)}"}
else
:ok
end
end
# Normalizes an entry to ensure consistent structure
defp normalize_entry(entry) do
%{
python_module: entry.python_module,
python_version: entry.python_version,
elixir_module: entry.elixir_module,
generated_at: normalize_datetime(entry.generated_at),
path: entry.path,
files: entry.files,
stats: normalize_stats(entry.stats)
}
end
# Normalizes stats to ensure atoms as keys
defp normalize_stats(stats) do
%{
functions: stats[:functions] || stats["functions"],
classes: stats[:classes] || stats["classes"],
submodules: stats[:submodules] || stats["submodules"]
}
end
# Normalizes datetime - accepts DateTime or string
defp normalize_datetime(%DateTime{} = dt), do: dt
defp normalize_datetime(string) when is_binary(string) do
case DateTime.from_iso8601(string) do
{:ok, dt, _offset} -> dt
{:error, _} -> raise ArgumentError, "Invalid datetime string: #{string}"
end
end
# Builds the registry data structure for JSON serialization
defp build_registry_data do
libraries =
registry_get(fn state ->
state
|> Enum.map(fn {name, entry} ->
{name, serialize_entry(entry)}
end)
|> Enum.into(%{})
end)
data = %{
"version" => @registry_version,
"generated_at" => DateTime.utc_now() |> DateTime.to_iso8601(),
"libraries" => libraries
}
{:ok, data}
end
# Serializes a registry entry for JSON
defp serialize_entry(entry) do
%{
"python_module" => entry.python_module,
"python_version" => entry.python_version,
"elixir_module" => entry.elixir_module,
"generated_at" => DateTime.to_iso8601(entry.generated_at),
"path" => entry.path,
"files" => entry.files,
"stats" => %{
"functions" => entry.stats.functions,
"classes" => entry.stats.classes,
"submodules" => entry.stats.submodules
}
}
end
# Parses registry data from JSON
defp parse_registry_data(%{"libraries" => libraries}) when is_map(libraries) do
parsed =
libraries
|> Enum.map(fn {name, entry} ->
case deserialize_entry(entry) do
{:ok, parsed_entry} -> {name, parsed_entry}
{:error, reason} -> {:error, {name, reason}}
end
end)
# Check for any errors
errors =
Enum.filter(parsed, fn
{:error, _} -> true
_ -> false
end)
if length(errors) > 0 do
{:error, "Invalid entries: #{inspect(errors)}"}
else
{:ok, Enum.into(parsed, %{})}
end
end
defp parse_registry_data(_data) do
{:error, "Invalid registry format: missing 'libraries' key"}
end
# Deserializes a registry entry from JSON
defp deserialize_entry(entry) when is_map(entry) do
{:ok, generated_at, _offset} = DateTime.from_iso8601(entry["generated_at"])
parsed = %{
python_module: entry["python_module"],
python_version: entry["python_version"],
elixir_module: entry["elixir_module"],
generated_at: generated_at,
path: entry["path"],
files: entry["files"],
stats: %{
functions: entry["stats"]["functions"],
classes: entry["stats"]["classes"],
submodules: entry["stats"]["submodules"]
}
}
{:ok, parsed}
rescue
e ->
{:error, "Failed to deserialize entry: #{inspect(e)}"}
end
defp deserialize_entry(_entry) do
{:error, "Entry must be a map"}
end
defp registry_get(fun) do
with_registry(fn -> Agent.get(__MODULE__, fun) end)
end
defp registry_update(fun) do
with_registry(fn -> Agent.update(__MODULE__, fun) end)
end
defp with_registry(fun) do
ensure_started()
try do
fun.()
catch
:exit, {:noproc, _} ->
start_registry()
fun.()
:exit, :noproc ->
start_registry()
fun.()
end
end
end
defmodule SnakeBridge.RuntimeClient do
@moduledoc """
Behaviour for runtime clients that execute SnakeBridge payloads.
The default runtime client is `Snakepit`, but tests can override
this via the `:runtime_client` config.
"""
@type tool :: String.t()
@type payload :: map()
@type opts :: keyword()
@type callback :: (term() -> any())
@callback execute(tool(), payload(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
@callback execute_stream(tool(), payload(), callback(), opts()) ::
:ok | {:error, Snakepit.Error.t()}
end
defmodule SnakeBridge.Runtime do
@moduledoc """
Thin payload helper for SnakeBridge that delegates execution to Snakepit.
This module is compile-time agnostic and focuses on building payloads that
match the Snakepit Prime runtime contract.
"""
alias SnakeBridge.SessionManager
alias SnakeBridge.Types
require Logger
@type module_ref :: module()
@type function_name :: atom() | String.t()
@type args :: list()
@type opts :: keyword()
@protocol_version 1
@min_supported_version 1
# Process dictionary key for auto-session
@auto_session_key :snakebridge_auto_session
@doc """
Call a Python function.
## Parameters
- `module` - Either a generated SnakeBridge module atom OR a Python module path string
- `function` - Function name (atom or string)
- `args` - Positional arguments (list)
- `opts` - Options including kwargs, :idempotent, :__runtime__
## Examples
# With generated module
{:ok, result} = SnakeBridge.Runtime.call(Numpy, :mean, [[1,2,3]])
# With string module path (dynamic)
{:ok, result} = SnakeBridge.Runtime.call("numpy", "mean", [[1,2,3]])
{:ok, result} = SnakeBridge.Runtime.call("math", :sqrt, [16])
"""
@spec call(module_ref() | String.t(), function_name() | String.t(), args(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def call(module, function, args \\ [], opts \\ [])
# String module path - delegate to dynamic
def call(module, function, args, opts) when is_binary(module) do
function_name = to_string(function)
call_dynamic(module, function_name, args, opts)
end
# Atom module - existing behavior
def call(module, function, args, opts) when is_atom(module) do
{kwargs, idempotent, extra_args, runtime_opts} = split_opts(opts)
encoded_args = encode_args(args ++ extra_args)
encoded_kwargs = encode_kwargs(kwargs)
# Determine session_id ONCE using correct priority
session_id = resolve_session_id(runtime_opts)
payload =
base_payload(module, function, encoded_args, encoded_kwargs, idempotent)
|> Map.put("session_id", session_id)
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = call_metadata(payload, module, function, "function")
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> apply_error_mode()
|> decode_result()
end
@doc """
Calls any Python function dynamically without requiring generated bindings.
This is the no-codegen escape hatch for calling functions that were not
scanned during compilation.
"""
@spec call_dynamic(String.t(), function_name(), args(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def call_dynamic(module_path, function, args \\ [], opts \\ []) when is_binary(module_path) do
{args, opts} = normalize_args_opts(args, opts)
{kwargs, idempotent, extra_args, runtime_opts} = split_opts(opts)
encoded_args = encode_args(args ++ extra_args)
encoded_kwargs = encode_kwargs(kwargs)
# Determine session_id ONCE - this is the single source of truth
session_id = resolve_session_id(runtime_opts)
library = module_path |> String.split(".") |> List.first()
payload =
protocol_payload()
|> Map.put("call_type", "dynamic")
|> Map.put("module_path", module_path)
|> Map.put("library", library)
|> Map.put("function", to_string(function))
|> Map.put("args", encoded_args)
|> Map.put("kwargs", encoded_kwargs)
|> Map.put("idempotent", idempotent)
|> maybe_put_session_id(session_id)
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = %{
module: module_path,
function: to_string(function),
library: library,
python_module: module_path,
call_type: "dynamic"
}
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> apply_error_mode()
|> decode_result()
end
@spec call_helper(String.t(), args(), opts() | map()) :: {:ok, term()} | {:error, term()}
def call_helper(helper, args \\ [], opts \\ [])
def call_helper(helper, args, opts) when is_map(opts) do
encoded_args = encode_args(args)
encoded_kwargs = encode_kwargs(stringify_keys(opts))
# Map opts cannot have __runtime__, use context/auto-session
session_id = resolve_session_id([])
payload =
helper_payload(helper, encoded_args, encoded_kwargs, false)
|> Map.put("session_id", session_id)
runtime_opts =
[]
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = helper_metadata(helper)
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> classify_helper_result(helper)
|> decode_result()
end
def call_helper(helper, args, opts) when is_list(opts) do
{kwargs, idempotent, extra_args, runtime_opts} = split_opts(opts)
encoded_args = encode_args(args ++ extra_args)
encoded_kwargs = encode_kwargs(kwargs)
# Determine session_id ONCE using correct priority
session_id = resolve_session_id(runtime_opts)
payload =
helper_payload(helper, encoded_args, encoded_kwargs, idempotent)
|> Map.put("session_id", session_id)
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = helper_metadata(helper)
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> classify_helper_result(helper)
|> decode_result()
end
@doc """
Stream results from a Python generator/iterator.
## Parameters
- `module` - Either a generated SnakeBridge module atom OR a Python module path string
- `function` - Function name (atom or string)
- `args` - Positional arguments (list)
- `opts` - Options including kwargs
- `callback` - Function called for each streamed item
## Performance
When called with a **generated module atom**, this function can use Snakepit's
native gRPC streaming for efficient data transfer.
When called with a **string module path**, this delegates to `stream_dynamic/5`
which uses RPC-per-item iteration. See `stream_dynamic/5` docs for performance
guidance on large streams.
## Examples
# With string module path (dynamic, RPC-per-item)
SnakeBridge.Runtime.stream("pandas", "read_csv", ["file.csv"], [chunksize: 100], fn chunk ->
process(chunk)
end)
# With generated module (native streaming when available)
SnakeBridge.Runtime.stream(MyApp.Pandas, :read_csv, ["file.csv"], [chunksize: 100], fn chunk ->
process(chunk)
end)
"""
@spec stream(module_ref() | String.t(), function_name() | String.t(), args(), opts(), (term() ->
any())) ::
:ok | {:ok, :done} | {:error, Snakepit.Error.t()}
def stream(module, function, args \\ [], opts \\ [], callback)
# String module path - use stream_dynamic
def stream(module, function, args, opts, callback)
when is_binary(module) and is_function(callback, 1) do
function_name = to_string(function)
stream_dynamic(module, function_name, args, opts, callback)
end
# Atom module - existing behavior
def stream(module, function, args, opts, callback)
when is_atom(module) and is_function(callback, 1) do
{kwargs, idempotent, extra_args, runtime_opts} = split_opts(opts)
encoded_args = encode_args(args ++ extra_args)
encoded_kwargs = encode_kwargs(kwargs)
# Determine session_id ONCE using correct priority
session_id = resolve_session_id(runtime_opts)
payload =
base_payload(module, function, encoded_args, encoded_kwargs, idempotent)
|> Map.put("session_id", session_id)
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :stream)
|> ensure_session_opt(session_id)
metadata = call_metadata(payload, module, function, "stream")
decode_callback = fn chunk -> callback.(Types.decode(chunk)) end
execute_with_telemetry(metadata, fn ->
runtime_client().execute_stream(
"snakebridge.stream",
payload,
decode_callback,
runtime_opts
)
end)
|> apply_error_mode()
end
@doc """
Stream results from a Python generator using dynamic dispatch.
Creates a stream reference and iterates via stream_next until exhausted.
## Performance Note
Dynamic streaming uses an RPC-per-item approach: each item from the Python
iterator triggers a separate `stream_next` gRPC call. This is correct and
safe but may be slow for large streams (thousands of items).
For high-throughput streaming workloads, consider:
- **Generated streaming wrappers**: Use `SnakeBridge.stream/5` with compiled
modules, which can leverage Snakepit's server-side streaming for better
throughput.
- **Batched iteration**: Have Python yield batches of items rather than
individual items.
- **Dedicated data transfer**: For very large datasets, consider writing
Python results to files/databases and loading from Elixir.
Dynamic streaming is ideal for convenience and moderate-sized iterables.
"""
@spec stream_dynamic(String.t(), String.t(), args(), opts(), (term() -> any())) ::
{:ok, :done} | {:error, term()}
def stream_dynamic(module_path, function, args, opts, callback)
when is_binary(module_path) and is_function(callback, 1) do
case call_dynamic(module_path, function, args, opts) do
{:ok, %SnakeBridge.StreamRef{} = stream_ref} ->
stream_iterate(stream_ref, callback, [])
{:ok, %SnakeBridge.Ref{} = ref} ->
# Try to iterate if it's an iterator
stream_iterate_ref(ref, callback, [])
{:ok, other} ->
# Not a stream/iterator - return as single value
callback.(other)
{:ok, :done}
error ->
error
end
end
defp stream_iterate(stream_ref, callback, opts) do
case stream_next(stream_ref, opts) do
{:ok, item} ->
callback.(item)
stream_iterate(stream_ref, callback, opts)
{:error, :stop_iteration} ->
{:ok, :done}
{:error, _} = error ->
error
end
end
defp stream_iterate_ref(ref, callback, opts) do
case call_method(ref, :__iter__, [], opts) do
{:ok, %SnakeBridge.StreamRef{} = stream_ref} ->
stream_iterate(stream_ref, callback, opts)
{:ok, %SnakeBridge.Ref{} = iter_ref} ->
stream_iterate_ref_next(iter_ref, callback, opts)
{:ok, other} ->
callback.(other)
{:ok, :done}
{:error, reason} ->
if stop_iteration?(reason) do
{:ok, :done}
else
stream_iterate_ref_next(ref, callback, opts)
end
end
end
defp stream_iterate_ref_next(ref, callback, opts) do
case call_method(ref, :__next__, [], opts) do
{:ok, item} ->
callback.(item)
stream_iterate_ref_next(ref, callback, opts)
{:error, reason} ->
if stop_iteration?(reason) do
{:ok, :done}
else
{:error, reason}
end
end
end
defp stop_iteration?(reason) when is_map(reason) do
type =
Map.get(reason, :python_type) || Map.get(reason, "python_type") ||
Map.get(reason, :error_type) || Map.get(reason, "error_type")
type == "StopIteration"
end
@doc """
Gets the next item from a Python iterator or generator.
Each call makes a separate RPC to Python. For high-throughput streaming,
see the performance note on `stream_dynamic/5`.
"""
@spec stream_next(SnakeBridge.StreamRef.t(), opts()) ::
{:ok, term()} | {:error, :stop_iteration} | {:error, Snakepit.Error.t()}
def stream_next(stream_ref, opts \\ []) do
{_args, opts} = normalize_args_opts([], opts)
{_, _, _, runtime_opts} = split_opts(opts)
wire_ref = SnakeBridge.StreamRef.to_wire_format(stream_ref)
# Single source of truth: prioritize runtime_opts, then stream_ref session, then context
session_id = resolve_session_id(runtime_opts, stream_ref)
library =
case stream_ref.library do
lib when is_binary(lib) and lib != "" -> lib
_ -> "unknown"
end
payload =
protocol_payload()
|> Map.put("call_type", "stream_next")
|> Map.put("stream_ref", wire_ref)
|> Map.put("library", library)
|> maybe_put_session_id(session_id)
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :stream)
|> ensure_session_opt(session_id)
result =
runtime_client().execute("snakebridge.call", payload, runtime_opts)
|> apply_error_mode()
case result do
{:ok, %{"__type__" => "stop_iteration"}} ->
{:error, :stop_iteration}
{:ok, value} ->
{:ok, Types.decode(value)}
error ->
error
end
end
@doc """
Gets the length of a Python iterable (if supported).
"""
@spec stream_len(SnakeBridge.StreamRef.t(), opts()) ::
{:ok, non_neg_integer()} | {:error, term()}
def stream_len(stream_ref, opts \\ []) do
wire_ref = SnakeBridge.StreamRef.to_wire_format(stream_ref)
call_method(wire_ref, :__len__, [], opts)
end
@spec call_class(module_ref(), function_name(), args(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def call_class(module, function, args \\ [], opts \\ []) do
{kwargs, idempotent, extra_args, runtime_opts} = split_opts(opts)
encoded_args = encode_args(args ++ extra_args)
encoded_kwargs = encode_kwargs(kwargs)
# Determine session_id ONCE using correct priority
session_id = resolve_session_id(runtime_opts)
payload =
module
|> base_payload(function, encoded_args, encoded_kwargs, idempotent)
|> Map.put("call_type", "class")
|> Map.put("class", python_class_name(module))
|> Map.put("session_id", session_id)
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = call_metadata(payload, module, function, "class")
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> apply_error_mode()
|> decode_result()
end
@spec call_method(SnakeBridge.Ref.t() | map(), function_name(), args(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def call_method(ref, function, args \\ [], opts \\ []) do
{kwargs, idempotent, extra_args, runtime_opts} = split_opts(opts)
encoded_args = encode_args(args ++ extra_args)
encoded_kwargs = encode_kwargs(kwargs)
wire_ref = normalize_ref(ref)
# Single source of truth: prioritize runtime_opts, then ref session, then context
session_id = resolve_session_id(runtime_opts, wire_ref)
payload =
wire_ref
|> base_payload_for_ref(function, encoded_args, encoded_kwargs, idempotent)
|> Map.put("call_type", "method")
|> Map.put("instance", wire_ref)
|> Map.put("session_id", session_id)
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = ref_metadata(payload, function, "method")
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> apply_error_mode()
|> decode_result()
end
@doc """
Retrieves a module-level attribute (constant, class, etc.).
## Parameters
- `module` - Either a generated SnakeBridge module atom OR a Python module path string
- `attr` - Attribute name (atom or string)
- `opts` - Runtime options
## Examples
# Get math.pi
{:ok, pi} = SnakeBridge.Runtime.get_module_attr("math", "pi")
{:ok, pi} = SnakeBridge.Runtime.get_module_attr("math", :pi)
"""
@spec get_module_attr(module_ref() | String.t(), atom() | String.t(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def get_module_attr(module, attr, opts \\ [])
# String module path
def get_module_attr(module, attr, opts) when is_binary(module) do
{_kwargs, _idempotent, _extra_args, runtime_opts} = split_opts(opts)
attr_name = to_string(attr)
# Determine session_id ONCE using correct priority
session_id = resolve_session_id(runtime_opts)
payload =
protocol_payload()
|> Map.put("call_type", "module_attr")
|> Map.put("python_module", module)
|> Map.put("library", library_from_module_path(module))
|> Map.put("attr", attr_name)
|> Map.put("session_id", session_id)
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = %{
module: module,
function: attr_name,
library: library_from_module_path(module),
python_module: module,
call_type: "module_attr"
}
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> apply_error_mode()
|> decode_result()
end
# Atom module - existing behavior
def get_module_attr(module, attr, opts) when is_atom(module) do
{kwargs, idempotent, _extra_args, runtime_opts} = split_opts(opts)
encoded_kwargs = encode_kwargs(kwargs)
# Determine session_id ONCE using correct priority
session_id = resolve_session_id(runtime_opts)
payload =
module
|> base_payload(attr, [], encoded_kwargs, idempotent)
|> Map.put("call_type", "module_attr")
|> Map.put("attr", to_string(attr))
|> Map.put("session_id", session_id)
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = call_metadata(payload, module, attr, "module_attr")
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> apply_error_mode()
|> decode_result()
end
# Helper to extract library name from module path
defp library_from_module_path(module_path) when is_binary(module_path) do
module_path
|> String.split(".")
|> List.first()
end
@spec get_attr(SnakeBridge.Ref.t(), atom() | String.t(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def get_attr(ref, attr, opts \\ []) do
{kwargs, idempotent, _extra_args, runtime_opts} = split_opts(opts)
encoded_kwargs = encode_kwargs(kwargs)
wire_ref = normalize_ref(ref)
# Single source of truth: prioritize runtime_opts, then ref session, then context
session_id = resolve_session_id(runtime_opts, wire_ref)
payload =
wire_ref
|> base_payload_for_ref(attr, [], encoded_kwargs, idempotent)
|> Map.put("call_type", "get_attr")
|> Map.put("instance", wire_ref)
|> Map.put("attr", to_string(attr))
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = ref_metadata(payload, attr, "get_attr")
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> apply_error_mode()
|> decode_result()
end
@doc false
def build_module_attr_payload(module, attr) do
module
|> base_payload(attr, [], %{}, false)
|> Map.put("call_type", "module_attr")
|> Map.put("attr", to_string(attr))
end
@doc false
def build_dynamic_payload(module_path, function, args, opts) do
{args, opts} = normalize_args_opts(args, opts)
{kwargs, idempotent, extra_args, _runtime_opts} = split_opts(opts)
protocol_payload()
|> Map.put("call_type", "dynamic")
|> Map.put("module_path", module_path)
|> Map.put("function", to_string(function))
|> Map.put("args", List.wrap(args ++ extra_args))
|> Map.put("kwargs", Map.new(kwargs, fn {key, value} -> {to_string(key), value} end))
|> Map.put("idempotent", idempotent)
|> maybe_put_session_id(current_session_id())
end
@spec set_attr(SnakeBridge.Ref.t(), atom() | String.t(), term(), opts()) ::
{:ok, term()} | {:error, Snakepit.Error.t()}
def set_attr(ref, attr, value, opts \\ []) do
{kwargs, idempotent, _extra_args, runtime_opts} = split_opts(opts)
encoded_kwargs = encode_kwargs(kwargs)
encoded_args = encode_args([value])
wire_ref = normalize_ref(ref)
# Single source of truth: prioritize runtime_opts, then ref session, then context
session_id = resolve_session_id(runtime_opts, wire_ref)
payload =
wire_ref
|> base_payload_for_ref(attr, encoded_args, encoded_kwargs, idempotent)
|> Map.put("call_type", "set_attr")
|> Map.put("instance", wire_ref)
|> Map.put("attr", to_string(attr))
runtime_opts =
runtime_opts
|> apply_runtime_defaults(payload, :call)
|> ensure_session_opt(session_id)
metadata = ref_metadata(payload, attr, "set_attr")
execute_with_telemetry(metadata, fn ->
runtime_client().execute("snakebridge.call", payload, runtime_opts)
end)
|> apply_error_mode()
|> decode_result()
end
@spec release_ref(SnakeBridge.Ref.t(), opts()) :: :ok | {:error, Snakepit.Error.t()}
def release_ref(ref, opts \\ []) do
{_, _, _, runtime_opts} = split_opts(opts)
wire_ref = normalize_ref(ref)
# Single source of truth: prioritize runtime_opts, then ref session, then context
session_id = resolve_session_id(runtime_opts, wire_ref)
runtime_opts = ensure_session_opt(runtime_opts, session_id)
payload =
protocol_payload()
|> Map.put("ref", wire_ref)
|> maybe_put_session_id(session_id)
runtime_client().execute("snakebridge.release_ref", payload, runtime_opts)
|> apply_error_mode()
|> normalize_release_result()
end
@spec release_session(String.t(), opts()) :: :ok | {:error, Snakepit.Error.t()}
def release_session(session_id, opts \\ []) when is_binary(session_id) do
{_, _, _, runtime_opts} = split_opts(opts)
runtime_opts = ensure_session_opt(runtime_opts, session_id)
payload =
protocol_payload()
|> Map.put("session_id", session_id)
runtime_client().execute("snakebridge.release_session", payload, runtime_opts)
|> apply_error_mode()
|> normalize_release_result()
end
defp runtime_client do
Application.get_env(:snakebridge, :runtime_client, Snakepit)
end
defp execute_with_telemetry(metadata, fun) do
start_time = System.monotonic_time()
emit_runtime_event(
[:snakepit, :python, :call, :start],
%{system_time: System.system_time()},
metadata
)
try do
result = fun.()
case result do
{:error, reason} ->
emit_runtime_event(
[:snakepit, :python, :call, :exception],
%{duration: System.monotonic_time() - start_time},
Map.put(metadata, :error, reason)
)
_ ->
emit_runtime_event(
[:snakepit, :python, :call, :stop],
%{duration: System.monotonic_time() - start_time},
metadata
)
end
result
rescue
exception ->
emit_runtime_event(
[:snakepit, :python, :call, :exception],
%{duration: System.monotonic_time() - start_time},
Map.put(metadata, :reason, exception)
)
reraise exception, __STACKTRACE__
end
end
defp emit_runtime_event(event, measurements, metadata) do
case Application.ensure_all_started(:telemetry) do
{:ok, _} -> :telemetry.execute(event, measurements, metadata)
{:error, _} -> :ok
end
end
defp call_metadata(payload, module, function, call_type) do
%{
module: module,
function: to_string(function),
library: payload["library"],
python_module: payload["python_module"],
call_type: call_type
}
end
defp ref_metadata(payload, function, call_type) do
%{
module: payload["python_module"],
function: to_string(function),
library: payload["library"],
python_module: payload["python_module"],
call_type: call_type
}
end
defp helper_metadata(helper) do
%{
module: helper,
function: helper,
library: helper_library(helper),
python_module: helper_library(helper),
call_type: "helper"
}
end
defp split_opts(opts) do
extra_args = Keyword.get(opts, :__args__, [])
idempotent = Keyword.get(opts, :idempotent, false)
runtime_opts = Keyword.get(opts, :__runtime__, [])
kwargs =
opts
|> Keyword.drop([:__args__, :idempotent, :__runtime__])
|> Enum.into(%{}, fn {key, value} -> {to_string(key), value} end)
{kwargs, idempotent, List.wrap(extra_args), runtime_opts}
end
# Session ID single source of truth: determine once, use everywhere
# Priority: runtime_opts override > ref session_id > context session > auto-session
@doc false
def resolve_session_id(runtime_opts, ref \\ nil) do
session_id_from_runtime_opts(runtime_opts) ||
session_id_from_ref(ref) ||
current_session_id()
end
defp session_id_from_runtime_opts(runtime_opts) when is_list(runtime_opts) do
Keyword.get(runtime_opts, :session_id)
end
defp session_id_from_runtime_opts(_), do: nil
defp session_id_from_ref(%SnakeBridge.Ref{session_id: id}) when is_binary(id), do: id
defp session_id_from_ref(%SnakeBridge.StreamRef{session_id: id}) when is_binary(id), do: id
defp session_id_from_ref(ref) when is_map(ref) do
if Map.has_key?(ref, "session_id") or Map.has_key?(ref, :session_id) do
ref_field(ref, "session_id")
end
end
defp session_id_from_ref(_), do: nil
defp ensure_session_opt(runtime_opts, session_id) when is_binary(session_id) do
cond do
runtime_opts == nil ->
[session_id: session_id]
is_list(runtime_opts) ->
Keyword.put_new(runtime_opts, :session_id, session_id)
true ->
runtime_opts
end
end
defp ensure_session_opt(runtime_opts, _session_id), do: runtime_opts
# ============================================================================
# Runtime Timeout Defaults
# ============================================================================
# Applies runtime timeout defaults to the given runtime options.
#
# This function merges profile-based timeouts with user-provided options,
# respecting the following priority (highest to lowest):
# 1. Explicit user-provided options (e.g., `timeout: 60_000`)
# 2. User-selected profile (via `timeout_profile:` or `profile:`)
# 3. Library-specific profile (from `runtime.library_profiles` config)
# 4. Default profile based on call kind (`:default` for calls, `:streaming` for streams)
@doc false
@spec apply_runtime_defaults(keyword() | nil, map(), atom()) :: keyword()
def apply_runtime_defaults(runtime_opts, payload, call_kind) do
runtime_opts = List.wrap(runtime_opts || [])
library = payload["library"]
profile = resolve_timeout_profile(runtime_opts, library, call_kind)
profile_opts = get_profile_opts(profile)
# Merge: profile defaults < user overrides
merged =
profile_opts
|> Keyword.merge(runtime_opts)
merged
|> Keyword.put_new(:timeout_profile, profile)
|> Keyword.put_new(:timeout, SnakeBridge.Defaults.runtime_default_timeout())
|> maybe_put_stream_defaults(call_kind)
end
defp resolve_timeout_profile(runtime_opts, library, call_kind) do
# Priority: explicit > library_profiles > global default
Keyword.get(runtime_opts, :timeout_profile) ||
Keyword.get(runtime_opts, :profile) ||
library_profile(library) ||
SnakeBridge.Defaults.runtime_timeout_profile(call_kind)
end
defp get_profile_opts(profile) do
SnakeBridge.Defaults.runtime_profiles()
|> Map.get(profile, [])
end
defp library_profile(nil), do: nil
defp library_profile(library) when is_binary(library) do
profiles = SnakeBridge.Defaults.runtime_library_profiles()
Map.get(profiles, library) ||
Map.get(profiles, String.to_existing_atom(library))
rescue
ArgumentError -> nil
end
defp library_profile(_), do: nil
defp maybe_put_stream_defaults(opts, :stream) do
Keyword.put_new(opts, :stream_timeout, SnakeBridge.Defaults.runtime_default_stream_timeout())
end
defp maybe_put_stream_defaults(opts, _), do: opts
@doc false
@spec normalize_args_opts(list(), keyword()) :: {list(), keyword()}
def normalize_args_opts(args, opts) do
if opts == [] and Keyword.keyword?(args) do
{[], args}
else
{args, opts}
end
end
defp base_payload(module, function, args, kwargs, idempotent) do
python_module = python_module_name(module)
%{
"protocol_version" => @protocol_version,
"min_supported_version" => @min_supported_version,
"library" => library_name(module, python_module),
"python_module" => python_module,
"function" => to_string(function),
"args" => List.wrap(args),
"kwargs" => kwargs,
"idempotent" => idempotent
}
|> maybe_put_session_id(current_session_id())
end
defp base_payload_for_ref(ref, function, args, kwargs, idempotent) do
python_module =
ref_field(ref, "python_module") || ref_field(ref, "library") || python_module_name(ref)
library = ref_field(ref, "library") || library_name(ref, python_module)
session_id = ref_field(ref, "session_id") || current_session_id()
%{
"protocol_version" => @protocol_version,
"min_supported_version" => @min_supported_version,
"library" => library,
"python_module" => python_module,
"function" => to_string(function),
"args" => List.wrap(args),
"kwargs" => kwargs,
"idempotent" => idempotent
}
|> maybe_put_session_id(session_id)
end
defp python_module_name(module) when is_atom(module) do
if function_exported?(module, :__snakebridge_python_name__, 0) do
module.__snakebridge_python_name__()
else
module
|> Module.split()
|> Enum.map_join(".", &Macro.underscore/1)
end
end
defp python_module_name(%{python_module: python_module}) when is_binary(python_module),
do: python_module
defp python_module_name(_), do: "unknown"
defp library_name(module, python_module) when is_atom(module) do
if function_exported?(module, :__snakebridge_library__, 0) do
module.__snakebridge_library__()
else
python_module |> String.split(".") |> List.first()
end
end
defp library_name(_module, python_module) do
python_module |> String.split(".") |> List.first()
end
defp python_class_name(module) when is_atom(module) do
if function_exported?(module, :__snakebridge_python_class__, 0) do
module.__snakebridge_python_class__()
else
module |> Module.split() |> List.last()
end
end
defp helper_payload(helper, args, kwargs, idempotent) do
%{
"protocol_version" => @protocol_version,
"min_supported_version" => @min_supported_version,
"call_type" => "helper",
"helper" => helper,
"function" => helper,
"library" => helper_library(helper),
"args" => List.wrap(args),
"kwargs" => kwargs,
"idempotent" => idempotent,
"helper_config" => SnakeBridge.Helpers.payload_config(SnakeBridge.Helpers.runtime_config())
}
|> maybe_put_session_id(current_session_id())
end
@doc false
def protocol_payload do
%{
"protocol_version" => @protocol_version,
"min_supported_version" => @min_supported_version
}
end
defp helper_library(helper) when is_binary(helper) do
case String.split(helper, ".", parts: 2) do
[library, _rest] -> library
_ -> "unknown"
end
end
defp helper_library(_), do: "unknown"
defp current_session_id do
case SnakeBridge.SessionContext.current() do
%{session_id: session_id} when is_binary(session_id) -> session_id
_ -> ensure_auto_session()
end
end
# Auto-session management
@doc """
Returns the current session ID (explicit or auto-generated).
This is useful for debugging or when you need to know which session is active.
"""
@spec current_session() :: String.t()
def current_session do
current_session_id()
end
@doc """
Clears the auto-session for the current process.
Useful for testing or when you want to force a new session.
Does NOT release the session on the Python side - use `release_auto_session/0` for that.
"""
@spec clear_auto_session() :: :ok
def clear_auto_session do
Process.delete(@auto_session_key)
:ok
end
@doc """
Releases and clears the auto-session for the current process.
This releases all refs associated with the session on both Elixir and Python sides.
"""
@spec release_auto_session() :: :ok
def release_auto_session do
case Process.get(@auto_session_key) do
nil ->
:ok
session_id ->
# Release on Python side
release_session(session_id)
# Unregister from SessionManager
SessionManager.unregister_session(session_id)
# Clear from process dictionary
Process.delete(@auto_session_key)
:ok
end
end
defp ensure_auto_session do
case Process.get(@auto_session_key) do
nil ->
session_id = generate_auto_session_id()
setup_auto_session(session_id)
session_id
session_id ->
session_id
end
end
defp generate_auto_session_id do
pid_string = self() |> :erlang.pid_to_list() |> to_string()
timestamp = System.system_time(:millisecond)
"auto_#{pid_string}_#{timestamp}"
end
defp setup_auto_session(session_id) do
# Store in process dictionary
Process.put(@auto_session_key, session_id)
# Register with SessionManager for monitoring
# This ensures cleanup when the process dies
SessionManager.register_session(session_id, self())
# Ensure Snakepit session exists (if SessionStore is available)
ensure_snakepit_session(session_id)
end
defp ensure_snakepit_session(session_id) do
# Only call if SessionStore module is available
# Use apply/3 to avoid compile-time warnings about undefined module
if Code.ensure_loaded?(Snakepit.SessionStore) do
# credo:disable-for-next-line Credo.Check.Refactor.Apply
case apply(Snakepit.SessionStore, :create_session, [session_id]) do
{:ok, _} ->
:ok
{:error, :already_exists} ->
:ok
{:error, reason} ->
Logger.warning("Failed to create Snakepit session #{session_id}: #{inspect(reason)}")
:ok
end
else
:ok
end
end
defp maybe_put_session_id(payload, nil), do: payload
defp maybe_put_session_id(payload, session_id) when is_binary(session_id) do
Map.put(payload, "session_id", session_id)
end
defp classify_helper_result({:error, reason}, helper) do
{:error, classify_helper_error(reason, helper)}
end
defp classify_helper_result(result, _helper), do: result
defp classify_helper_error({:invalid_parameter, :json_encode_failed, message}, _helper) do
SnakeBridge.SerializationError.new(message)
end
defp classify_helper_error(
%{python_type: "SnakeBridgeHelperNotFoundError", message: message},
helper
) do
helper_name = extract_helper_name(message) || helper
SnakeBridge.HelperNotFoundError.new(helper_name)
end
defp classify_helper_error(
%{python_type: "SnakeBridgeSerializationError", message: message},
_helper
) do
SnakeBridge.SerializationError.new(message)
end
defp classify_helper_error(reason, _helper), do: reason
defp extract_helper_name(message) when is_binary(message) do
case Regex.run(~r/Helper ['"]([^'"]+)['"]/, message) do
[_, helper] -> helper
_ -> nil
end
end
defp extract_helper_name(_), do: nil
defp encode_args(args) do
args
|> List.wrap()
|> Enum.map(&Types.encode/1)
end
defp encode_kwargs(kwargs) do
kwargs
|> Enum.into(%{}, fn {key, value} -> {to_string(key), Types.encode(value)} end)
end
defp stringify_keys(map) when is_map(map) do
Enum.into(map, %{}, fn {key, value} -> {to_string(key), value} end)
end
defp decode_result({:ok, value}), do: {:ok, Types.decode(value)}
defp decode_result(result), do: result
defp apply_error_mode({:error, reason}) do
case error_mode() do
:raw ->
{:error, reason}
:translated ->
{:error, translate_reason(reason)}
:raise_translated ->
translated = translate_reason(reason)
if translated == reason do
{:error, reason}
else
raise translated
end
end
end
defp apply_error_mode(result), do: result
defp normalize_release_result({:ok, _}), do: :ok
defp normalize_release_result(:ok), do: :ok
defp normalize_release_result(result), do: result
defp translate_reason(reason) do
case python_error_payload(reason) do
{_message, traceback, type} when is_binary(type) ->
translated = SnakeBridge.ErrorTranslator.translate(reason, traceback)
if translated == reason, do: reason, else: translated
{message, traceback, _type} when is_binary(message) ->
translated =
SnakeBridge.ErrorTranslator.translate(%RuntimeError{message: message}, traceback)
case translated do
%RuntimeError{} -> reason
_ -> translated
end
_ ->
reason
end
end
defp python_error_payload(error) when is_map(error) do
{extract_error_message(error), extract_error_traceback(error), extract_error_type(error)}
end
defp python_error_payload(_), do: {nil, nil, nil}
defp extract_error_message(error) do
get_first_present(error, [:message, "message", :error, "error"])
end
defp extract_error_traceback(error) do
get_first_present(error, [:traceback, "traceback", :python_traceback, "python_traceback"])
end
defp extract_error_type(error) do
get_first_present(error, [:python_type, "python_type", :error_type, "error_type"])
end
defp get_first_present(map, keys) do
Enum.find_value(keys, fn key -> Map.get(map, key) end)
end
defp error_mode do
Application.get_env(:snakebridge, :error_mode, :raw)
end
defp ref_field(ref, "python_module") when is_map(ref),
do: Map.get(ref, "python_module") || Map.get(ref, :python_module)
defp ref_field(ref, "library") when is_map(ref),
do: Map.get(ref, "library") || Map.get(ref, :library)
defp ref_field(ref, "session_id") when is_map(ref),
do: Map.get(ref, "session_id") || Map.get(ref, :session_id)
defp ref_field(ref, "id") when is_map(ref),
do: Map.get(ref, "id") || Map.get(ref, :id) || Map.get(ref, "ref_id") || Map.get(ref, :ref_id)
defp ref_field(_ref, _key), do: nil
defp normalize_ref(%SnakeBridge.Ref{} = ref), do: SnakeBridge.Ref.to_wire_format(ref)
defp normalize_ref(ref) when is_map(ref) do
if Map.get(ref, "__type__") == "ref" or Map.get(ref, :__type__) == "ref" do
SnakeBridge.Ref.to_wire_format(ref)
else
ref
end
end
defp normalize_ref(ref), do: ref
end
defmodule SnakeBridge.ScanError do
@moduledoc """
Structured error for scan failures.
"""
defexception [:failures]
@type t :: %__MODULE__{failures: list(map())}
@impl Exception
def message(%__MODULE__{failures: failures}) do
failures
|> Enum.map_join("\n", fn %{path: path, reason: reason} ->
" - #{path}: #{inspect(reason)}"
end)
|> then(&("Scan failed for #{length(failures)} file(s):\n" <> &1))
end
end
defmodule SnakeBridge.Scanner do
@moduledoc """
Scans project source files for Python library calls.
"""
@type call_ref :: {module(), atom(), non_neg_integer()}
@spec scan_project(SnakeBridge.Config.t()) :: [call_ref()]
def scan_project(config) do
start_time = System.monotonic_time()
library_modules = Enum.map(config.libraries, & &1.module_name)
files = source_files(config)
{calls, failures} =
files
|> Task.async_stream(&scan_file(&1, library_modules))
|> Enum.zip(files)
|> Enum.reduce({[], []}, fn
{{:ok, calls}, _path}, {acc_calls, acc_failures} ->
{calls ++ acc_calls, acc_failures}
{{:exit, reason}, path}, {acc_calls, acc_failures} ->
{acc_calls, [%{path: path, reason: reason, type: :exit} | acc_failures]}
{{:error, reason}, path}, {acc_calls, acc_failures} ->
{acc_calls, [%{path: path, reason: reason, type: :error} | acc_failures]}
end)
calls =
calls
|> Enum.uniq()
|> Enum.sort()
SnakeBridge.Telemetry.scan_stop(
start_time,
length(files),
length(calls),
config.scan_paths || ["lib"]
)
if failures != [] do
raise SnakeBridge.ScanError, failures: Enum.reverse(failures)
end
calls
end
defp source_files(config) do
scan_paths = config.scan_paths || ["lib"]
scan_exclude = config.scan_exclude || []
scan_paths
|> Enum.flat_map(&Path.wildcard(Path.join(&1, "**/*.ex")))
|> Enum.reject(fn path ->
in_generated_dir?(path, config.generated_dir) or excluded_path?(path, scan_exclude)
end)
end
defp in_generated_dir?(path, generated_dir) do
String.starts_with?(path, generated_dir)
end
defp excluded_path?(path, patterns) do
Enum.any?(patterns, fn pattern ->
path in Path.wildcard(pattern)
end)
end
defp scan_file(path, library_modules) do
case File.read(path) do
{:ok, content} ->
case Code.string_to_quoted(content, file: path) do
{:ok, ast} ->
context = build_context(ast, library_modules)
extract_calls(ast, context)
{:error, _} ->
[]
end
{:error, _} ->
[]
end
end
defp build_context(ast, library_modules) do
{_, context} =
Macro.prewalk(ast, %{aliases: %{}, imports: []}, fn
{:alias, _, [{:__aliases__, _, parts} | opts]}, ctx ->
module = Module.concat(parts)
if library_module?(module, library_modules) do
alias_name = alias_name(parts, opts)
{nil, put_in(ctx, [:aliases, alias_name], module)}
else
{nil, ctx}
end
{:import, _, [{:__aliases__, _, parts} | opts]}, ctx ->
module = Module.concat(parts)
if library_module?(module, library_modules) do
{nil, update_in(ctx, [:imports], &[{module, opts} | &1])}
else
{nil, ctx}
end
node, ctx ->
{node, ctx}
end)
Map.put(context, :library_modules, library_modules)
end
defp alias_name(parts, opts) do
case Keyword.get(opts, :as) do
{:__aliases__, _, [name]} -> name
nil -> List.last(parts)
end
end
defp extract_calls(ast, context) do
{_, calls} =
Macro.prewalk(ast, [], fn
{{:., _, [{:__aliases__, _, parts}, function]}, _, args} = node, acc
when is_atom(function) and is_list(args) ->
module = resolve_module(parts, context)
if module do
{node, [{module, function, length(args)} | acc]}
else
{node, acc}
end
{function, _, args} = node, acc
when is_atom(function) and is_list(args) ->
case find_import(function, length(args), context) do
{:ok, module} -> {node, [{module, function, length(args)} | acc]}
:not_found -> {node, acc}
end
node, acc ->
{node, acc}
end)
calls
end
defp resolve_module(parts, context) do
module = Module.concat(parts)
case parts do
[name] when is_atom(name) ->
Map.get(context.aliases, name) ||
if library_module?(module, context.library_modules), do: module, else: nil
_ ->
if library_module?(module, context.library_modules), do: module, else: nil
end
end
defp library_module?(module, library_modules) do
module_parts = Module.split(module)
Enum.any?(library_modules, fn library_module ->
library_parts = Module.split(library_module)
Enum.take(module_parts, length(library_parts)) == library_parts
end)
end
defp find_import(function, arity, context) do
Enum.find_value(context.imports, :not_found, fn {module, opts} ->
only = Keyword.get(opts, :only, nil)
except = Keyword.get(opts, :except, [])
cond do
{function, arity} in except -> nil
only && {function, arity} not in only -> nil
true -> {:ok, module}
end
end)
end
end
defmodule SnakeBridge.SerializationError do
@moduledoc """
Raised when attempting to encode a value that cannot be serialized for Python.
SnakeBridge supports encoding:
- Primitives: `nil`, booleans, integers, floats, strings
- Collections: lists, maps, tuples, MapSets
- Special types: atoms, DateTime, Date, Time, SnakeBridge.Bytes
- References: SnakeBridge.Ref, SnakeBridge.StreamRef
- Functions: anonymous functions (as callbacks)
- Special floats: `:infinity`, `:neg_infinity`, `:nan`
Types that cannot be serialized:
- PIDs, ports, references
- Custom structs without serialization support
- File handles, sockets, other system resources
## Resolution
For unsupported types, you have several options:
1. **Create a Python object and pass the ref**:
{:ok, ref} = SnakeBridge.call("module", "create_object", [...])
SnakeBridge.call("module", "use_object", [ref])
2. **Convert to a supported type**:
# Instead of passing a PID
SnakeBridge.call("module", "fn", [inspect(pid)])
# Or extract relevant data
SnakeBridge.call("module", "fn", [pid_to_list(pid)])
3. **Use explicit bytes for binary data**:
SnakeBridge.call("module", "fn", [SnakeBridge.bytes(binary)])
"""
defexception [:message, :value, :type]
@type t :: %__MODULE__{
message: String.t(),
value: term(),
type: atom() | module()
}
@impl true
def exception(opts) when is_list(opts) do
value = Keyword.fetch!(opts, :value)
type = get_type(value)
message = build_message(value, type)
%__MODULE__{
message: message,
value: value,
type: type
}
end
@doc """
Creates a SerializationError from a message string.
This is used for error messages from the Python side.
"""
@spec new(String.t() | nil) :: t()
def new(message \\ nil) do
%__MODULE__{
message: message || "Arguments are not JSON-serializable",
value: nil,
type: :unknown
}
end
defp get_type(value) when is_pid(value), do: :pid
defp get_type(value) when is_port(value), do: :port
defp get_type(value) when is_reference(value), do: :reference
defp get_type(%{__struct__: struct_name}), do: struct_name
defp get_type(_), do: :unknown
defp build_message(value, type) do
"""
Cannot serialize value of type #{inspect(type)} for Python.
Value: #{inspect(value, limit: 50, printable_limit: 100)}
SnakeBridge cannot automatically serialize this type. See the module documentation
for SnakeBridge.SerializationError for resolution options.
"""
end
end
defmodule SnakeBridge.SessionContext do
@moduledoc """
Provides scoped session context for Python calls.
Sessions control the lifecycle of Python object references (refs). Each session
is isolated, meaning refs from one session cannot be used in another.
## Automatic vs Explicit Sessions
By default, SnakeBridge creates an auto-session for each Elixir process. This is
convenient for most use cases where Python objects don't need to be shared.
Use explicit sessions when you need:
- Multiple processes to access the same Python objects
- Long-lived refs that outlive a single request/task
- Fine-grained control over cleanup timing
## Usage
# Explicit session with custom ID
SnakeBridge.SessionContext.with_session([session_id: "my-session"], fn ->
{:ok, model} = SnakeBridge.call("sklearn.linear_model", "LinearRegression", [])
# model ref is accessible by other processes using "my-session"
model
end)
# Simple scoped session (auto-generated ID)
SnakeBridge.SessionContext.with_session(fn ->
# All Python calls here use the same session
{:ok, df} = SnakeBridge.call("pandas", "DataFrame", [[1, 2, 3]])
{:ok, mean} = SnakeBridge.method(df, "mean", [])
mean
end)
## Session Cleanup
Sessions are automatically cleaned up when:
- The owning process dies (auto-sessions)
- `SnakeBridge.Runtime.release_session/1` is called explicitly
- Refs exceed TTL (SessionContext default: 1 hour) or max count (default 10,000)
## Sharing Refs Across Processes
To share Python objects across processes, use the same explicit session_id:
# Process A
session_id = "shared-#{System.unique_integer()}"
SessionContext.with_session([session_id: session_id], fn ->
{:ok, ref} = SnakeBridge.call("heavy_model", "load", [])
send(process_b, {:model, session_id, ref})
end)
# Process B - can use the ref if it adopts the same session
receive do
{:model, session_id, ref} ->
SessionContext.with_session([session_id: session_id], fn ->
{:ok, result} = SnakeBridge.method(ref, "predict", [data])
result
end)
end
## Options
- `:session_id` - Custom session ID (default: auto-generated)
- `:max_refs` - Maximum refs per session (default: 10,000)
- `:ttl_seconds` - Session time-to-live in seconds (default: 3600, i.e., 1 hour)
- `:tags` - Custom metadata for debugging
"""
alias Snakepit.Bridge.SessionStore
@context_key :snakebridge_session_context
defstruct [
:session_id,
:owner_pid,
:created_at,
max_refs: 10_000,
ttl_seconds: 3600,
tags: %{}
]
@type t :: %__MODULE__{
session_id: String.t(),
owner_pid: pid(),
created_at: integer(),
max_refs: pos_integer(),
ttl_seconds: pos_integer(),
tags: map()
}
@doc """
Creates a new session context.
"""
@spec create(keyword()) :: t()
def create(opts \\ []) do
default_max_refs = Application.get_env(:snakebridge, :session_max_refs, 10_000)
default_ttl = Application.get_env(:snakebridge, :session_ttl_seconds, 3600)
%__MODULE__{
session_id: Keyword.get(opts, :session_id, generate_session_id()),
owner_pid: Keyword.get(opts, :owner_pid, self()),
created_at: System.system_time(:second),
max_refs: Keyword.get(opts, :max_refs, default_max_refs),
ttl_seconds: Keyword.get(opts, :ttl_seconds, default_ttl),
tags: Keyword.get(opts, :tags, %{})
}
end
@doc """
Gets the current session context from the process dictionary.
"""
@spec current() :: t() | nil
def current do
Process.get(@context_key)
end
@doc """
Sets the current session context in the process dictionary.
"""
@spec put_current(t()) :: t() | nil
def put_current(context) do
Process.put(@context_key, context)
end
@doc """
Clears the current session context.
"""
@spec clear_current() :: t() | nil
def clear_current do
Process.delete(@context_key)
end
@doc """
Executes a function within a session context.
The session is automatically registered and will be released
when the owner process dies.
"""
@spec with_session((-> result)) :: result when result: term()
def with_session(fun) when is_function(fun, 0) do
with_session([], fun)
end
@spec with_session(keyword(), (-> result)) :: result when result: term()
def with_session(opts, fun) when is_list(opts) and is_function(fun, 0) do
context = create(opts)
case SnakeBridge.SessionManager.register_session(context.session_id, context.owner_pid) do
:ok -> :ok
{:error, :already_exists} -> :ok
end
ensure_snakepit_session(context.session_id)
old_context = put_current(context)
try do
fun.()
after
if old_context do
put_current(old_context)
else
clear_current()
end
end
end
defp generate_session_id do
"session_#{:erlang.unique_integer([:positive])}_#{System.system_time(:millisecond)}"
end
defp ensure_snakepit_session(session_id) when is_binary(session_id) do
if Code.ensure_loaded?(SessionStore) and Process.whereis(SessionStore) do
_ = SessionStore.create_session(session_id)
end
:ok
end
end
defmodule SnakeBridge.SessionManager do
@moduledoc """
Manages Python session lifecycle with process monitoring.
Sessions are automatically released when their owner process dies,
preventing memory leaks in long-running applications.
"""
use GenServer
require Logger
@type session_id :: String.t()
@type ref :: map()
# Client API
def start_link(opts \\ []) do
GenServer.start_link(__MODULE__, opts, name: __MODULE__)
end
@doc """
Registers a new session with an owner process.
The session will be released when the owner dies.
"""
@spec register_session(session_id(), pid()) :: :ok | {:error, :already_exists}
def register_session(session_id, owner_pid) do
GenServer.call(__MODULE__, {:register_session, session_id, owner_pid})
end
@doc """
Registers a ref with its session for tracking.
"""
@spec register_ref(session_id(), ref()) :: :ok | {:error, :session_not_found}
def register_ref(session_id, ref) do
GenServer.call(__MODULE__, {:register_ref, session_id, ref})
end
@doc """
Checks if a session exists.
"""
@spec session_exists?(session_id()) :: boolean()
def session_exists?(session_id) do
GenServer.call(__MODULE__, {:session_exists?, session_id})
end
@doc """
Lists all refs in a session.
"""
@spec list_refs(session_id()) :: [ref()]
def list_refs(session_id) do
GenServer.call(__MODULE__, {:list_refs, session_id})
end
@doc """
Explicitly releases a session and all its refs.
"""
@spec release_session(session_id()) :: :ok
def release_session(session_id) do
GenServer.call(__MODULE__, {:release_session, session_id})
end
@doc """
Unregisters a session without releasing refs on the Python side.
Typically called when manually cleaning up before process death,
or when the caller has already released the session.
"""
@spec unregister_session(session_id()) :: :ok
def unregister_session(session_id) do
GenServer.call(__MODULE__, {:unregister_session, session_id})
end
# Server Implementation
@impl true
def init(_opts) do
state = %{
# session_id => %{owner_pid, monitor_ref, refs, created_at}
sessions: %{},
# monitor_ref => session_id
monitors: %{}
}
{:ok, state}
end
@impl true
def handle_call({:register_session, session_id, owner_pid}, _from, state) do
if Map.has_key?(state.sessions, session_id) do
{:reply, {:error, :already_exists}, state}
else
monitor_ref = Process.monitor(owner_pid)
session_data = %{
owner_pid: owner_pid,
monitor_ref: monitor_ref,
refs: [],
created_at: System.system_time(:second)
}
new_state = %{
state
| sessions: Map.put(state.sessions, session_id, session_data),
monitors: Map.put(state.monitors, monitor_ref, session_id)
}
{:reply, :ok, new_state}
end
end
@impl true
def handle_call({:register_ref, session_id, ref}, _from, state) do
case Map.get(state.sessions, session_id) do
nil ->
{:reply, {:error, :session_not_found}, state}
session_data ->
updated = %{session_data | refs: [ref | session_data.refs]}
new_state = put_in(state.sessions[session_id], updated)
{:reply, :ok, new_state}
end
end
@impl true
def handle_call({:session_exists?, session_id}, _from, state) do
{:reply, Map.has_key?(state.sessions, session_id), state}
end
@impl true
def handle_call({:list_refs, session_id}, _from, state) do
refs =
case Map.get(state.sessions, session_id) do
nil -> []
session_data -> session_data.refs
end
{:reply, refs, state}
end
@impl true
def handle_call({:release_session, session_id}, _from, state) do
new_state = do_release_session(state, session_id)
{:reply, :ok, new_state}
end
@impl true
def handle_call({:unregister_session, session_id}, _from, state) do
case Map.get(state.sessions, session_id) do
nil ->
{:reply, :ok, state}
%{monitor_ref: ref} ->
Process.demonitor(ref, [:flush])
new_state = %{
state
| sessions: Map.delete(state.sessions, session_id),
monitors: Map.delete(state.monitors, ref)
}
{:reply, :ok, new_state}
end
end
@impl true
def handle_info({:DOWN, monitor_ref, :process, _pid, _reason}, state) do
case Map.get(state.monitors, monitor_ref) do
nil ->
{:noreply, state}
session_id ->
Logger.debug("Session owner died, releasing session: #{session_id}")
new_state = do_release_session(state, session_id)
{:noreply, new_state}
end
end
defp do_release_session(state, session_id) do
case Map.get(state.sessions, session_id) do
nil ->
state
session_data ->
Process.demonitor(session_data.monitor_ref, [:flush])
Task.start(fn ->
try do
SnakeBridge.Runtime.release_session(session_id, [])
rescue
_ -> :ok
catch
:exit, _ -> :ok
end
end)
%{
state
| sessions: Map.delete(state.sessions, session_id),
monitors: Map.delete(state.monitors, session_data.monitor_ref)
}
end
end
end
defmodule SnakeBridge.SessionMismatchError do
@moduledoc """
Raised when a ref is used with a different session than it was created in.
SnakeBridge refs are session-scoped: a ref created in session A cannot be
used in session B. This error indicates a ref is being used across session
boundaries.
## Fields
- `:ref_id` - The ref ID that caused the mismatch
- `:expected_session` - The session ID the ref belongs to
- `:actual_session` - The session ID the ref was used in
- `:message` - Human-readable error message
"""
defexception [:ref_id, :expected_session, :actual_session, :message]
@type t :: %__MODULE__{
ref_id: String.t() | nil,
expected_session: String.t() | nil,
actual_session: String.t() | nil,
message: String.t()
}
@impl Exception
def exception(opts) when is_list(opts) do
ref_id = Keyword.get(opts, :ref_id)
expected_session = Keyword.get(opts, :expected_session)
actual_session = Keyword.get(opts, :actual_session)
message =
Keyword.get(opts, :message) || build_message(ref_id, expected_session, actual_session)
%__MODULE__{
ref_id: ref_id,
expected_session: expected_session,
actual_session: actual_session,
message: message
}
end
@impl Exception
def message(%__MODULE__{message: message}), do: message
defp build_message(ref_id, expected, actual) do
"SnakeBridge reference '#{ref_id || "unknown"}' belongs to session '#{expected || "unknown"}' " <>
"but was used in session '#{actual || "unknown"}'. Refs cannot be shared across sessions."
end
end
if Code.ensure_loaded?(Snakepit.PyRef) == false do
defmodule Snakepit.PyRef do
@moduledoc """
Reference to a Python object managed by Snakepit.
This is a stub type definition used when the Snakepit library is not loaded.
When Snakepit is available, its actual `Snakepit.PyRef` module takes precedence.
"""
@type t :: SnakeBridge.Ref.t()
end
end
if Code.ensure_loaded?(Snakepit.ZeroCopyRef) == false do
defmodule Snakepit.ZeroCopyRef do
@moduledoc """
Reference to a zero-copy Python buffer managed by Snakepit.
This is a stub type definition used when the Snakepit library is not loaded.
When Snakepit is available, its actual `Snakepit.ZeroCopyRef` module takes precedence.
"""
@type t :: term()
end
end
if Code.ensure_loaded?(Snakepit.Error) == false do
defmodule Snakepit.Error do
@moduledoc """
Error struct for Snakepit operations.
This is a stub type definition used when the Snakepit library is not loaded.
When Snakepit is available, its actual `Snakepit.Error` module takes precedence.
"""
@type t :: term()
@doc "Creates a validation error."
def validation_error(message, metadata \\ %{}) do
%{type: :validation_error, message: message, metadata: metadata}
end
end
end
defmodule SnakeBridge.StreamRef do
@moduledoc """
Represents a Python iterator or generator as an Elixir stream.
Implements the `Enumerable` protocol for lazy iteration.
"""
defstruct [
:ref_id,
:session_id,
:stream_type,
:python_module,
:library,
exhausted: false
]
@type t :: %__MODULE__{
ref_id: String.t(),
session_id: String.t(),
stream_type: String.t(),
python_module: String.t(),
library: String.t(),
exhausted: boolean()
}
@doc """
Creates a StreamRef from a decoded wire format.
"""
@spec from_wire_format(map()) :: t()
def from_wire_format(map) when is_map(map) do
%__MODULE__{
ref_id: map["id"],
session_id: map["session_id"],
stream_type: map["stream_type"] || "iterator",
python_module: map["python_module"],
library: map["library"],
exhausted: false
}
end
@doc """
Converts back to wire format for Python calls.
"""
@spec to_wire_format(t()) :: map()
def to_wire_format(%__MODULE__{} = ref) do
%{
"__type__" => "ref",
"id" => ref.ref_id,
"session_id" => ref.session_id,
"python_module" => ref.python_module,
"library" => ref.library
}
end
end
defimpl Enumerable, for: SnakeBridge.StreamRef do
alias SnakeBridge.{Runtime, StreamRef}
def count(%StreamRef{stream_type: "generator"}), do: {:error, __MODULE__}
def count(%StreamRef{} = ref) do
case Runtime.stream_len(ref) do
{:ok, len} when is_integer(len) -> {:ok, len}
_ -> {:error, __MODULE__}
end
end
def member?(%StreamRef{}, _value), do: {:error, __MODULE__}
def slice(%StreamRef{}), do: {:error, __MODULE__}
def reduce(%StreamRef{exhausted: true}, {:cont, acc}, _fun) do
{:done, acc}
end
def reduce(%StreamRef{} = ref, {:cont, acc}, fun) do
case Runtime.stream_next(ref) do
{:ok, value} ->
reduce(ref, fun.(value, acc), fun)
{:error, :stop_iteration} ->
{:done, acc}
{:error, reason} ->
{:halted, {:error, reason}}
end
end
def reduce(_ref, {:halt, acc}, _fun), do: {:halted, acc}
def reduce(ref, {:suspend, acc}, fun), do: {:suspended, acc, &reduce(ref, &1, fun)}
end
defmodule SnakeBridge.Telemetry do
@moduledoc """
Telemetry event definitions for SnakeBridge.
This module provides instrumentation for compile-time operations including:
- Source scanning
- Python introspection
- Code generation
- Lock file verification
## Event List
| Event | Measurements | Metadata |
|-------|-------------|----------|
| `[:snakebridge, :compile, :start]` | `system_time` | `libraries`, `strict` |
| `[:snakebridge, :compile, :stop]` | `duration`, `symbols_generated`, `files_written` | `libraries`, `mode` |
| `[:snakebridge, :compile, :exception]` | `duration` | `reason`, `stacktrace` |
| `[:snakebridge, :compile, :scan, :stop]` | `duration`, `files_scanned`, `symbols_found` | `library`, `phase`, `details` |
| `[:snakebridge, :compile, :introspect, :start]` | `system_time` | `library`, `phase`, `details` |
| `[:snakebridge, :compile, :introspect, :stop]` | `duration`, `symbols_introspected`, `cache_hits` | `library`, `phase`, `details` |
| `[:snakebridge, :compile, :generate, :stop]` | `duration`, `bytes_written`, `functions_generated`, `classes_generated` | `library`, `phase`, `details` |
| `[:snakebridge, :docs, :fetch]` | `duration` | `module`, `function`, `source` |
| `[:snakebridge, :lock, :verify]` | `duration` | `result`, `warnings` |
## Usage
# Attach handlers in your application
SnakeBridge.Telemetry.Handlers.Logger.attach()
# Compile-time events are automatically emitted during mix compile
"""
# ============================================================
# COMPILATION EVENTS
# ============================================================
@doc """
Emits compile start event.
## Measurements
- `system_time` - System time when compilation started
## Metadata
- `library` - `:all`
- `phase` - `:compile`
- `details` - `%{libraries: [...], strict: boolean()}`
"""
@spec compile_start([atom()], boolean()) :: :ok
def compile_start(libraries, strict) do
emit(
[:snakebridge, :compile, :start],
%{system_time: System.system_time()},
%{library: :all, phase: :compile, details: %{libraries: libraries, strict: strict}}
)
end
@doc """
Emits compile stop event.
## Measurements
- `duration` - Time in native units
- `symbols_generated` - Number of symbols generated
- `files_written` - Number of files written
## Metadata
- `library` - `:all`
- `phase` - `:compile`
- `details` - `%{libraries: [...], mode: :normal | :strict}`
"""
@spec compile_stop(integer(), non_neg_integer(), non_neg_integer(), [atom()], :normal | :strict) ::
:ok
def compile_stop(start_time, symbols, files, libraries, mode) do
emit(
[:snakebridge, :compile, :stop],
%{
duration: System.monotonic_time() - start_time,
symbols_generated: symbols,
files_written: files
},
%{library: :all, phase: :compile, details: %{libraries: libraries, mode: mode}}
)
end
@doc """
Emits compile exception event.
## Measurements
- `duration` - Time in native units
## Metadata
- `library` - `:all`
- `phase` - `:compile`
- `details` - `%{reason: term(), stacktrace: list()}`
"""
@spec compile_exception(integer(), term(), list()) :: :ok
def compile_exception(start_time, reason, stacktrace) do
emit(
[:snakebridge, :compile, :exception],
%{duration: System.monotonic_time() - start_time},
%{library: :all, phase: :compile, details: %{reason: reason, stacktrace: stacktrace}}
)
end
# ============================================================
# SCANNING EVENTS
# ============================================================
@doc """
Emits scan stop event.
## Measurements
- `duration` - Time in native units
- `files_scanned` - Number of files scanned
- `symbols_found` - Number of symbols found
## Metadata
- `library` - `:all`
- `phase` - `:scan`
- `details` - `%{paths: [String.t()]}`
"""
@spec scan_stop(integer(), non_neg_integer(), non_neg_integer(), [String.t()]) :: :ok
def scan_stop(start_time, files, symbols, paths) do
emit(
[:snakebridge, :compile, :scan, :stop],
%{
duration: System.monotonic_time() - start_time,
files_scanned: files,
symbols_found: symbols
},
%{library: :all, phase: :scan, details: %{paths: paths}}
)
end
# ============================================================
# INTROSPECTION EVENTS
# ============================================================
@doc """
Emits introspect start event.
## Measurements
- `system_time` - System time when introspection started
## Metadata
- `library` - Library atom being introspected
- `phase` - `:introspect`
- `details` - `%{batch_size: non_neg_integer()}`
"""
@spec introspect_start(atom(), non_neg_integer()) :: :ok
def introspect_start(library, batch_size) do
emit(
[:snakebridge, :compile, :introspect, :start],
%{system_time: System.system_time()},
%{library: library, phase: :introspect, details: %{batch_size: batch_size}}
)
end
@doc """
Emits introspect stop event.
## Measurements
- `duration` - Time in native units
- `symbols_introspected` - Number of symbols introspected
- `cache_hits` - Number of cache hits
## Metadata
- `library` - Library atom introspected
- `phase` - `:introspect`
- `details` - `%{python_time: integer()}`
"""
@spec introspect_stop(integer(), atom(), non_neg_integer(), non_neg_integer(), integer()) :: :ok
def introspect_stop(start_time, library, symbols, cache_hits, python_time) do
emit(
[:snakebridge, :compile, :introspect, :stop],
%{
duration: System.monotonic_time() - start_time,
symbols_introspected: symbols,
cache_hits: cache_hits
},
%{library: library, phase: :introspect, details: %{python_time: python_time}}
)
end
# ============================================================
# GENERATION EVENTS
# ============================================================
@doc """
Emits generate stop event.
## Measurements
- `duration` - Time in native units
- `bytes_written` - Number of bytes written
- `functions_generated` - Number of functions generated
- `classes_generated` - Number of classes generated
## Metadata
- `library` - Library atom generated
- `phase` - `:generate`
- `details` - `%{file: String.t()}`
"""
@spec generate_stop(
integer(),
atom(),
String.t(),
non_neg_integer(),
non_neg_integer(),
non_neg_integer()
) :: :ok
def generate_stop(start_time, library, file, bytes, functions, classes) do
emit(
[:snakebridge, :compile, :generate, :stop],
%{
duration: System.monotonic_time() - start_time,
bytes_written: bytes,
functions_generated: functions,
classes_generated: classes
},
%{library: library, phase: :generate, details: %{file: file}}
)
end
# ============================================================
# DOCUMENTATION EVENTS
# ============================================================
@doc """
Emits docs fetch event.
## Measurements
- `duration` - Time in native units
## Metadata
- `module` - Module fetched
- `function` - Function name
- `source` - `:cache`, `:python`, or `:metadata`
"""
@spec docs_fetch(integer(), module(), atom(), :cache | :python | :metadata) :: :ok
def docs_fetch(start_time, module, function, source) do
emit(
[:snakebridge, :docs, :fetch],
%{duration: System.monotonic_time() - start_time},
%{module: module, function: function, source: source}
)
end
# ============================================================
# LOCK FILE EVENTS
# ============================================================
@doc """
Emits lock verify event.
## Measurements
- `duration` - Time in native units
## Metadata
- `result` - `:ok`, `:warning`, or `:error`
- `warnings` - List of warning strings
"""
@spec lock_verify(integer(), :ok | :warning | :error, [String.t()]) :: :ok
def lock_verify(start_time, result, warnings \\ []) do
emit(
[:snakebridge, :lock, :verify],
%{duration: System.monotonic_time() - start_time},
%{result: result, warnings: warnings}
)
end
@doc """
Returns the expected metadata fields for an event.
"""
@spec event_metadata_schema([atom()]) :: [atom()]
def event_metadata_schema([:snakebridge, :compile | _]) do
[:library, :phase, :details]
end
def event_metadata_schema([:snakebridge, :runtime | _]) do
[:library, :function, :call_type]
end
def event_metadata_schema(_event), do: []
defp emit(event, measurements, metadata) do
case Application.ensure_all_started(:telemetry) do
{:ok, _} -> :telemetry.execute(event, measurements, metadata)
{:error, _} -> :ok
end
end
end
defmodule SnakeBridge.Types do
@moduledoc """
Public API for encoding and decoding Elixir types for Python interop.
This module provides a unified interface for type conversion between Elixir
and Python. It handles the serialization of Elixir-specific types (tuples,
MapSets, DateTime, etc.) into JSON-compatible formats and vice versa.
## Usage
# Encoding Elixir to JSON-compatible format
iex> SnakeBridge.Types.encode({:ok, 42})
%{
"__type__" => "tuple",
"__schema__" => 1,
"elements" => [%{"__type__" => "atom", "__schema__" => 1, "value" => "ok"}, 42]
}
# Decoding JSON-compatible format back to Elixir
iex> SnakeBridge.Types.decode(%{
...> "__type__" => "tuple",
...> "__schema__" => 1,
...> "elements" => [%{"__type__" => "atom", "__schema__" => 1, "value" => "ok"}, 42]
...> })
{:ok, 42}
## Type System
The type system uses tagged JSON representations to preserve type information
across the Elixir-Python boundary. See `SnakeBridge.Types.Encoder` and
`SnakeBridge.Types.Decoder` for details on supported types and their
representations.
## Round-trip Safety
All encoded values can be round-tripped (atoms depend on the decode allowlist):
iex> data = {:ok, MapSet.new([1, 2, 3])}
iex> data |> SnakeBridge.Types.encode() |> SnakeBridge.Types.decode()
{:ok, MapSet.new([1, 2, 3])}
"""
alias SnakeBridge.Types.{Decoder, Encoder}
@schema_version 1
@doc """
Returns the current SnakeBridge wire schema version for tagged values.
"""
@spec schema_version() :: pos_integer()
def schema_version, do: @schema_version
@doc """
Encodes an Elixir value into a JSON-compatible structure.
Delegates to `SnakeBridge.Types.Encoder.encode/1`.
## Examples
iex> SnakeBridge.Types.encode(:ok)
%{"__type__" => "atom", "__schema__" => 1, "value" => "ok"}
iex> SnakeBridge.Types.encode({:ok, 42})
%{
"__type__" => "tuple",
"__schema__" => 1,
"elements" => [%{"__type__" => "atom", "__schema__" => 1, "value" => "ok"}, 42]
}
iex> SnakeBridge.Types.encode(MapSet.new([1, 2, 3]))
%{"__type__" => "set", "__schema__" => 1, "elements" => [1, 2, 3]}
"""
@spec encode(term()) :: term()
defdelegate encode(value), to: Encoder
@doc """
Decodes a JSON-compatible structure back into Elixir types.
Delegates to `SnakeBridge.Types.Decoder.decode/1`.
## Examples
iex> SnakeBridge.Types.decode("ok")
"ok"
iex> SnakeBridge.Types.decode(%{
...> "__type__" => "tuple",
...> "__schema__" => 1,
...> "elements" => [%{"__type__" => "atom", "__schema__" => 1, "value" => "ok"}, 42]
...> })
{:ok, 42}
iex> SnakeBridge.Types.decode(%{"__type__" => "set", "elements" => [1, 2, 3]})
MapSet.new([1, 2, 3])
"""
@spec decode(term()) :: term()
defdelegate decode(value), to: Decoder
end
defmodule SnakeBridge.WheelConfig do
@moduledoc """
Configuration-based wheel variant selection.
"""
@default_config_path Path.join(["config", "wheel_variants.json"])
@doc """
Loads wheel configuration from file or uses defaults.
"""
@spec load_config() :: map()
def load_config do
case File.read(config_path()) do
{:ok, content} ->
Jason.decode!(content)
{:error, _} ->
default_config()
end
end
@doc """
Gets available variants for a package.
"""
@spec get_variants(String.t()) :: [String.t()]
def get_variants(package) do
config = load_config()
get_in(config, ["packages", package, "variants"]) || ["cpu"]
end
@doc """
Returns the configured packages that define variants.
"""
@spec packages() :: [String.t()]
def packages do
config = load_config()
config
|> Map.get("packages", %{})
|> Map.keys()
end
@doc """
Gets CUDA mapping for a version string.
"""
@spec get_cuda_mapping(String.t() | nil) :: String.t() | nil
def get_cuda_mapping(nil), do: nil
def get_cuda_mapping(version) do
config = load_config()
Map.get(config["cuda_mappings"] || %{}, version) ||
Map.get(config["cuda_mappings"] || %{}, normalize_cuda_version(version))
end
@doc """
Returns the configured ROCm variant, if any.
"""
@spec rocm_variant() :: String.t() | nil
def rocm_variant do
config = load_config()
config["rocm_variant"]
end
@doc false
def config_path do
Application.get_env(:snakebridge, :wheel_config_path) ||
Path.join(File.cwd!(), @default_config_path)
end
defp default_config do
%{
"packages" => %{
"torch" => %{"variants" => ["cpu", "cu118", "cu121", "cu124", "rocm5.7"]},
"torchvision" => %{"variants" => ["cpu", "cu118", "cu121", "cu124", "rocm5.7"]},
"torchaudio" => %{"variants" => ["cpu", "cu118", "cu121", "cu124", "rocm5.7"]}
},
"cuda_mappings" => %{
"11.7" => "cu118",
"11.8" => "cu118",
"12.0" => "cu121",
"12.1" => "cu121",
"12.2" => "cu121",
"12.3" => "cu124",
"12.4" => "cu124",
"12.5" => "cu124"
},
"rocm_variant" => "rocm5.7"
}
end
defp normalize_cuda_version(version) when is_binary(version) do
version
|> String.split(".")
|> Enum.take(2)
|> Enum.join()
end
defp normalize_cuda_version(_), do: nil
end
defmodule SnakeBridge.WheelSelector do
@moduledoc """
Selects the appropriate wheel variant for Python packages based on hardware.
PyTorch and related packages (torchvision, torchaudio) have different wheel
variants for different hardware configurations:
- `cpu` - CPU-only build
- `cu118` - CUDA 11.8
- `cu121` - CUDA 12.1
- `cu124` - CUDA 12.4
- `rocm5.7` - AMD ROCm 5.7
This module detects the current hardware and selects the appropriate variant.
## Examples
# Get the PyTorch variant for current hardware
variant = SnakeBridge.WheelSelector.pytorch_variant()
#=> "cu121" or "cpu"
# Get the index URL for pip
url = SnakeBridge.WheelSelector.pytorch_index_url()
#=> "https://download.pytorch.org/whl/cu121"
# Generate pip install command
cmd = SnakeBridge.WheelSelector.pip_install_command("torch", "2.1.0")
#=> "pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cu121"
"""
@type wheel_info :: %{
package: String.t(),
version: String.t(),
variant: String.t() | nil,
index_url: String.t() | nil
}
@doc """
Returns the PyTorch wheel variant for the current hardware.
## Examples
SnakeBridge.WheelSelector.pytorch_variant()
#=> "cu121" # On CUDA 12.1 system
#=> "cpu" # On CPU-only system
"""
@spec pytorch_variant() :: String.t()
def pytorch_variant do
caps = hardware_module().capabilities()
strategy_module().variant_for("torch", caps) || "cpu"
end
@doc """
Returns the PyTorch index URL for pip based on current hardware.
## Examples
SnakeBridge.WheelSelector.pytorch_index_url()
#=> "https://download.pytorch.org/whl/cu121"
"""
@spec pytorch_index_url() :: String.t()
def pytorch_index_url do
strategy_module().index_url_for_variant(pytorch_variant())
end
@doc """
Generates a pip install command for a package.
For PyTorch packages (torch, torchvision, torchaudio), includes the
appropriate --index-url for hardware-specific wheels.
## Examples
SnakeBridge.WheelSelector.pip_install_command("torch", "2.1.0")
#=> "pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cu121"
SnakeBridge.WheelSelector.pip_install_command("numpy", "1.26.4")
#=> "pip install numpy==1.26.4"
"""
@spec pip_install_command(String.t(), String.t()) :: String.t()
def pip_install_command(package, version) do
base = "pip install #{package}==#{version}"
wheel_info = select_wheel(package, version)
if wheel_info.index_url do
"#{base} --index-url #{wheel_info.index_url}"
else
base
end
end
@doc """
Normalizes a CUDA version string for wheel naming.
## Examples
SnakeBridge.WheelSelector.normalize_cuda_version("12.1")
#=> "121"
SnakeBridge.WheelSelector.normalize_cuda_version("11.8")
#=> "118"
"""
@spec normalize_cuda_version(String.t() | nil) :: String.t() | nil
def normalize_cuda_version(nil), do: nil
def normalize_cuda_version(version) when is_binary(version) do
version
|> String.split(".")
|> Enum.take(2)
|> Enum.join()
end
@doc """
Selects the appropriate wheel for a package based on current hardware.
Returns wheel info including variant and index URL if applicable.
## Examples
SnakeBridge.WheelSelector.select_wheel("torch", "2.1.0")
#=> %{package: "torch", version: "2.1.0", variant: "cu121", index_url: "..."}
SnakeBridge.WheelSelector.select_wheel("numpy", "1.26.4")
#=> %{package: "numpy", version: "1.26.4", variant: nil, index_url: nil}
"""
@spec select_wheel(String.t(), String.t()) :: wheel_info()
def select_wheel(package, version) do
strategy_module().select_wheel(package, version, hardware_module().capabilities())
end
@doc """
Checks if a package is a PyTorch package that needs hardware-specific wheels.
"""
@spec pytorch_package?(String.t()) :: boolean()
def pytorch_package?(package) do
package in SnakeBridge.WheelConfig.packages()
end
@doc """
Returns all available PyTorch variants for the given CUDA versions.
Useful for generating lock files that support multiple hardware configurations.
"""
@spec available_variants() :: [String.t()]
def available_variants do
available_variants("torch")
end
@spec available_variants(String.t()) :: [String.t()]
def available_variants(package) do
strategy_module().available_variants(package)
end
@doc """
Returns the best matching CUDA variant for a given CUDA version.
Falls back to the closest available version.
## Examples
SnakeBridge.WheelSelector.best_cuda_variant("12.3")
#=> "cu124"
SnakeBridge.WheelSelector.best_cuda_variant("12.1")
#=> "cu121"
"""
@spec best_cuda_variant(String.t() | nil) :: String.t()
def best_cuda_variant(cuda_version) do
strategy_module().best_cuda_variant(cuda_version)
end
# Private functions
defp hardware_module do
Application.get_env(:snakebridge, :hardware_module, Snakepit.Hardware)
end
defp strategy_module do
Application.get_env(:snakebridge, :wheel_strategy, SnakeBridge.WheelSelector.ConfigStrategy)
end
end
defmodule SnakeBridge.WithContext do
@moduledoc """
Provides Python context manager support via `with_python/2` macro.
Ensures `__exit__` is always called, even on exception.
## Example
SnakeBridge.with_python(file_ref) do
SnakeBridge.Dynamic.call(file_ref, :read, [])
end
"""
alias SnakeBridge.Runtime
@doc """
Executes a block with a Python context manager.
Calls `__enter__` before the block and guarantees `__exit__` after,
even if an exception occurs.
"""
defmacro with_python(ref, do: block) do
quote do
ref = unquote(ref)
case SnakeBridge.WithContext.call_enter(ref) do
{:ok, context_value} ->
var!(context) = context_value
_ = var!(context)
SnakeBridge.WithContext.execute_with_exit(ref, fn -> unquote(block) end)
{:error, reason} ->
{:error, reason}
end
end
end
@doc false
def execute_with_exit(ref, fun) when is_function(fun, 0) do
outcome =
try do
{:ok, fun.()}
rescue
exception ->
{:exception, exception, __STACKTRACE__}
end
case outcome do
{:ok, result} ->
call_exit(ref, nil)
result
{:exception, exception, stacktrace} ->
call_exit(ref, exception)
reraise exception, stacktrace
end
end
@doc """
Calls __enter__ on a Python context manager.
"""
@spec call_enter(SnakeBridge.Ref.t() | map(), keyword()) :: {:ok, term()} | {:error, term()}
def call_enter(ref, opts \\ []) do
Runtime.call_method(ref, :__enter__, [], opts)
end
@doc """
Calls __exit__ on a Python context manager.
"""
@spec call_exit(SnakeBridge.Ref.t() | map(), Exception.t() | nil, keyword()) ::
{:ok, term()} | {:error, term()}
def call_exit(ref, exception, opts \\ []) do
{exc_type, exc_value, exc_tb} =
if exception do
{
to_string(exception.__struct__),
Exception.message(exception),
nil
}
else
{nil, nil, nil}
end
Runtime.call_method(ref, :__exit__, [exc_type, exc_value, exc_tb], opts)
end
@doc false
def build_enter_payload(ref) do
wire_ref = SnakeBridge.Ref.to_wire_format(ref)
%{
"call_type" => "method",
"instance" => wire_ref,
"method" => "__enter__",
"args" => []
}
end
@doc false
def build_exit_payload(ref, exception) do
{exc_type, exc_value, exc_tb} =
if exception do
{to_string(exception.__struct__), Exception.message(exception), nil}
else
{nil, nil, nil}
end
wire_ref = SnakeBridge.Ref.to_wire_format(ref)
%{
"call_type" => "method",
"instance" => wire_ref,
"method" => "__exit__",
"args" => [exc_type, exc_value, exc_tb]
}
end
end
defmodule SnakeBridge do
@moduledoc """
Universal FFI bridge to Python.
SnakeBridge provides two ways to call Python:
1. **Generated wrappers** (compile-time): Type-safe, documented Elixir modules
generated from Python library introspection.
2. **Dynamic calls** (runtime): Direct calls to any Python module without
code generation, using string module paths.
## Universal FFI API
The universal FFI requires no code generation:
# Call any Python function
{:ok, result} = SnakeBridge.call("math", "sqrt", [16])
# Get module attributes
{:ok, pi} = SnakeBridge.get("math", "pi")
# Work with Python objects
{:ok, path} = SnakeBridge.call("pathlib", "Path", ["/tmp"])
{:ok, exists?} = SnakeBridge.method(path, "exists", [])
## Sessions and Ref Lifecycle
SnakeBridge automatically manages Python object sessions. Each Elixir process
gets an isolated session, and refs are automatically cleaned up when the
process terminates.
### Key Rules
1. **Refs are session-scoped**: A ref is only valid within its session. Don't
pass refs between processes without ensuring they share a session.
2. **Process death triggers cleanup**: When an Elixir process dies, its session
is released and all associated Python objects are garbage collected.
3. **Auto-session per process**: By default, each process gets an auto-session
(prefixed with `auto_`). Refs created in one process cannot be used from
another without explicit session sharing.
4. **Explicit sessions for sharing**: Use `SessionContext.with_session/2` with
a shared `session_id` to allow multiple processes to access the same refs.
5. **Ref TTL**: Python ref TTL is disabled by default. Enable via
`SNAKEBRIDGE_REF_TTL_SECONDS` environment variable. When enabled, refs
not accessed within the TTL window are cleaned up automatically.
6. **Max refs limit**: Each session can hold up to 10,000 refs by default.
Excess refs are pruned oldest-first. Configure via `SNAKEBRIDGE_REF_MAX`.
### Recommended Patterns
# Pattern 1: Single process, automatic cleanup
def process_data do
{:ok, df} = SnakeBridge.call("pandas", "read_csv", ["data.csv"])
{:ok, result} = SnakeBridge.method(df, "mean", [])
result # df is cleaned up when this process exits
end
# Pattern 2: Explicit session for long-lived refs
def with_shared_session(session_id) do
SnakeBridge.SessionContext.with_session([session_id: session_id], fn ->
{:ok, model} = SnakeBridge.call("sklearn.linear_model", "LinearRegression", [])
# Model ref can be accessed by other processes using same session_id
model
end)
end
# Pattern 3: Release refs explicitly when done
{:ok, ref} = SnakeBridge.call("io", "StringIO", ["test"])
# ... use ref ...
SnakeBridge.release_ref(ref) # Explicit cleanup
For explicit session control, use `SnakeBridge.SessionContext.with_session/1`.
## Type Mapping
| Elixir | Python |
|--------|--------|
| `nil` | `None` |
| `true`/`false` | `True`/`False` |
| integers | `int` |
| floats | `float` |
| strings | `str` |
| `SnakeBridge.bytes(data)` | `bytes` |
| lists | `list` |
| maps | `dict` |
| tuples | `tuple` |
| `MapSet` | `set` |
| atoms | tagged atom (decoded to string by default) |
| `DateTime` | `datetime` |
| `SnakeBridge.Ref` | Python object reference |
## Advanced Features (Opt-In)
SnakeBridge includes optional compile-time features that are disabled by default:
### Strict Mode
Enables compile-time verification of lock files and binding consistency.
Enable via `config :snakebridge, strict: true` or `SNAKEBRIDGE_STRICT=1`.
### Lock File Verification
Run `mix snakebridge.verify` to check that your lock file matches the current
environment. Useful in CI/CD to catch hardware/package drift.
### Wheel Selection
`SnakeBridge.WheelSelector` provides hardware-aware PyTorch wheel selection.
Call `WheelSelector.pytorch_variant/0` to get the appropriate CUDA/CPU variant.
### Helper Packs
Built-in helpers are enabled by default. Disable with:
config :snakebridge, helper_pack_enabled: false
### Environment Variables
| Variable | Default | Description |
|----------|---------|-------------|
| `SNAKEBRIDGE_STRICT` | `false` | Enable strict mode |
| `SNAKEBRIDGE_VERBOSE` | `false` | Verbose logging |
| `SNAKEBRIDGE_REF_TTL_SECONDS` | `0` | Ref TTL in seconds (0 = disabled) |
| `SNAKEBRIDGE_REF_MAX` | `10000` | Max refs per session |
| `SNAKEBRIDGE_STRICT_MODE` | `false` | Python strict mode (warns on ref accumulation) |
| `SNAKEBRIDGE_STRICT_MODE_THRESHOLD` | `1000` | Strict mode warning threshold |
"""
require SnakeBridge.WithContext
alias SnakeBridge.{Bytes, Dynamic, Ref, Runtime}
# ============================================================================
# Universal FFI API
# ============================================================================
@doc """
Call a Python function.
Accepts either a generated SnakeBridge module or a Python module path string.
## Parameters
- `module` - A generated module atom (e.g., `Numpy`) or a module path string (e.g., `"numpy"`)
- `function` - Function name as atom or string
- `args` - List of positional arguments (default: `[]`)
- `opts` - Keyword arguments passed to Python, plus:
- `:idempotent` - Mark call as cacheable (default: `false`)
- `:__runtime__` - Pass-through options to Snakepit
## Examples
# Call stdlib function
{:ok, 4.0} = SnakeBridge.call("math", "sqrt", [16])
# With keyword arguments
{:ok, 3.14} = SnakeBridge.call("builtins", "round", [3.14159], ndigits: 2)
# Submodule
{:ok, path} = SnakeBridge.call("os.path", "join", ["/tmp", "file.txt"])
# Create objects
{:ok, ref} = SnakeBridge.call("pathlib", "Path", ["."])
## Return Values
- `{:ok, value}` - Decoded Elixir value for JSON-serializable results
- `{:ok, %SnakeBridge.Ref{}}` - Reference for non-serializable Python objects
- `{:error, reason}` - Error from Python
## Notes
- String module paths trigger dynamic dispatch (no codegen required)
- Sessions are automatic; refs are isolated per Elixir process
- Non-JSON-serializable returns are wrapped in refs for safe access
"""
@spec call(module() | String.t(), atom() | String.t(), list(), keyword()) ::
{:ok, term()} | {:error, term()}
defdelegate call(module, function, args \\ [], opts \\ []), to: Runtime
@doc """
Call a Python function, raising on error.
Same as `call/4` but raises on error instead of returning `{:error, reason}`.
## Examples
result = SnakeBridge.call!("math", "sqrt", [16])
# => 4.0
# Raises on error
SnakeBridge.call!("nonexistent_module", "fn", [])
# ** (Snakepit.Error) ...
"""
@spec call!(module() | String.t(), atom() | String.t(), list(), keyword()) :: term()
def call!(module, function, args \\ [], opts \\ []) do
case call(module, function, args, opts) do
{:ok, result} -> result
{:error, error} -> raise error
end
end
@doc """
Get a module-level attribute from Python.
Retrieves constants, classes, or any attribute from a Python module.
## Parameters
- `module` - A generated module atom or a module path string
- `attr` - Attribute name as atom or string
- `opts` - Runtime options
## Examples
# Module constant
{:ok, pi} = SnakeBridge.get("math", "pi")
# => {:ok, 3.141592653589793}
# Module-level class (returns ref)
{:ok, path_class} = SnakeBridge.get("pathlib", "Path")
# Nested attribute
{:ok, sep} = SnakeBridge.get("os", "sep")
"""
@spec get(module() | String.t(), atom() | String.t(), keyword()) ::
{:ok, term()} | {:error, term()}
defdelegate get(module, attr, opts \\ []), to: Runtime, as: :get_module_attr
@doc """
Get a module-level attribute, raising on error.
"""
@spec get!(module() | String.t(), atom() | String.t(), keyword()) :: term()
def get!(module, attr, opts \\ []) do
case get(module, attr, opts) do
{:ok, result} -> result
{:error, error} -> raise error
end
end
@doc """
Stream results from a Python generator or iterator.
Calls a Python function that returns an iterable and invokes the callback
for each element.
## Parameters
- `module` - Module atom or path string
- `function` - Function name
- `args` - Positional arguments
- `opts` - Keyword arguments for the Python function
- `callback` - Function called with each streamed element
## Examples
# Process file in chunks
SnakeBridge.stream("pandas", "read_csv", ["large.csv"], [chunksize: 1000], fn chunk ->
IO.puts("Processing chunk")
end)
# Iterate range
SnakeBridge.stream("builtins", "range", [10], [], fn i ->
IO.puts("Got: \#{i}")
end)
## Return Value
- `{:ok, :done}` - Iteration completed successfully (for string module paths)
- `:ok` - Iteration completed successfully (for atom modules)
- `{:error, reason}` - Error during iteration
"""
@spec stream(module() | String.t(), atom() | String.t(), list(), keyword(), (term() -> term())) ::
:ok | {:ok, :done} | {:error, term()}
defdelegate stream(module, function, args, opts, callback), to: Runtime
@doc """
Call a method on a Python object reference.
## Parameters
- `ref` - A `SnakeBridge.Ref` from a previous call
- `method` - Method name as atom or string
- `args` - Positional arguments (default: `[]`)
- `opts` - Keyword arguments
## Examples
{:ok, path} = SnakeBridge.call("pathlib", "Path", ["."])
{:ok, exists?} = SnakeBridge.method(path, "exists", [])
{:ok, resolved} = SnakeBridge.method(path, "resolve", [])
# With arguments
{:ok, child} = SnakeBridge.method(path, "joinpath", ["subdir", "file.txt"])
## Notes
This is equivalent to `SnakeBridge.Dynamic.call/4` but with a clearer name
for the universal FFI context.
"""
@spec method(Ref.t(), atom() | String.t(), list(), keyword()) ::
{:ok, term()} | {:error, term()}
defdelegate method(ref, method, args \\ [], opts \\ []), to: Dynamic, as: :call
@doc """
Call a method on a ref, raising on error.
"""
@spec method!(Ref.t(), atom() | String.t(), list(), keyword()) :: term()
def method!(ref, method, args \\ [], opts \\ []) do
case method(ref, method, args, opts) do
{:ok, result} -> result
{:error, error} -> raise error
end
end
@doc """
Get an attribute from a Python object reference.
## Parameters
- `ref` - A `SnakeBridge.Ref` from a previous call
- `attr` - Attribute name as atom or string
- `opts` - Runtime options
## Examples
{:ok, path} = SnakeBridge.call("pathlib", "Path", ["/tmp/file.txt"])
{:ok, name} = SnakeBridge.attr(path, "name")
# => {:ok, "file.txt"}
{:ok, parent} = SnakeBridge.attr(path, "parent")
# => {:ok, %SnakeBridge.Ref{...}} # parent is also a Path
"""
@spec attr(Ref.t(), atom() | String.t(), keyword()) ::
{:ok, term()} | {:error, term()}
defdelegate attr(ref, attr, opts \\ []), to: Dynamic, as: :get_attr
@doc """
Get an attribute from a ref, raising on error.
"""
@spec attr!(Ref.t(), atom() | String.t(), keyword()) :: term()
def attr!(ref, attr, opts \\ []) do
case attr(ref, attr, opts) do
{:ok, result} -> result
{:error, error} -> raise error
end
end
@doc """
Set an attribute on a Python object reference.
## Parameters
- `ref` - A `SnakeBridge.Ref` from a previous call
- `attr` - Attribute name as atom or string
- `value` - New value for the attribute
- `opts` - Runtime options
## Examples
{:ok, obj} = SnakeBridge.call("some_module", "SomeClass", [])
{:ok, _} = SnakeBridge.set_attr(obj, "property", "new_value")
"""
@spec set_attr(Ref.t(), atom() | String.t(), term(), keyword()) ::
{:ok, term()} | {:error, term()}
defdelegate set_attr(ref, attr, value, opts \\ []), to: Dynamic
# ============================================================================
# Type Helpers
# ============================================================================
@doc """
Create a Bytes wrapper for explicit binary data.
By default, SnakeBridge encodes UTF-8 valid strings as Python `str`.
Use this function to explicitly send data as Python `bytes`.
## Examples
# Crypto
{:ok, hash_ref} = SnakeBridge.call("hashlib", "md5", [SnakeBridge.bytes("abc")])
{:ok, hex} = SnakeBridge.method(hash_ref, "hexdigest", [])
# Binary protocols
{:ok, packed} = SnakeBridge.call("struct", "pack", [">I", 12345])
# Base64
{:ok, encoded} = SnakeBridge.call("base64", "b64encode", [SnakeBridge.bytes("hello")])
## When to Use
Python distinguishes `str` (text) from `bytes` (binary). Use `bytes/1` for:
- Cryptographic operations (hashlib, hmac, cryptography)
- Binary packing (struct)
- Base64 encoding
- Network protocols
- File I/O in binary mode
"""
@spec bytes(binary()) :: Bytes.t()
def bytes(data) when is_binary(data) do
Bytes.new(data)
end
# ============================================================================
# Session Management
# ============================================================================
@doc """
Get the current session ID.
Returns the session ID for the current Elixir process. Sessions are
automatically created on first Python call.
## Examples
session_id = SnakeBridge.current_session()
# => "auto_<0.123.0>_1703944800000"
# With explicit session
SnakeBridge.SessionContext.with_session(session_id: "my_session", fn ->
SnakeBridge.current_session()
end)
# => "my_session"
"""
@spec current_session() :: String.t()
defdelegate current_session(), to: Runtime
@doc """
Release and clear the auto-session for the current process.
Call this to eagerly release Python object refs when you're done with
Python calls, rather than waiting for process termination.
## Examples
{:ok, ref} = SnakeBridge.call("numpy", "array", [[1,2,3]])
# ... use ref ...
SnakeBridge.release_auto_session() # Clean up now
## Notes
- This releases all refs in the current process's auto-session
- A new session is created automatically on the next Python call
- Use `SessionContext.with_session/1` for more fine-grained control
"""
@spec release_auto_session() :: :ok
defdelegate release_auto_session(), to: Runtime
@doc """
Releases a Python object reference, freeing memory in the Python process.
Call this to explicitly release a ref when you're done with it, rather than
waiting for session cleanup or process termination.
## Parameters
- `ref` - A `SnakeBridge.Ref` to release
- `opts` - Runtime options (optional)
## Examples
{:ok, ref} = SnakeBridge.call("pathlib", "Path", ["/tmp"])
# ... use ref ...
:ok = SnakeBridge.release_ref(ref)
## Notes
- After release, the ref is invalid and should not be used
- Releasing an already-released ref is a no-op
- For bulk cleanup, use `release_session/1` instead
"""
@spec release_ref(Ref.t(), keyword()) :: :ok | {:error, term()}
defdelegate release_ref(ref, opts \\ []), to: Runtime
@doc """
Releases all Python object references associated with a session.
Use this for bulk cleanup of all refs in a session, rather than releasing
them individually.
## Parameters
- `session_id` - The session ID to release
- `opts` - Runtime options (optional)
## Examples
session_id = SnakeBridge.current_session()
# ... create many refs ...
:ok = SnakeBridge.release_session(session_id)
## Notes
- After release, all refs from that session are invalid
- The session can still be reused for new calls
- For auto-sessions, prefer `release_auto_session/0`
"""
@spec release_session(String.t(), keyword()) :: :ok | {:error, term()}
defdelegate release_session(session_id, opts \\ []), to: Runtime
# ============================================================================
# Ref Utilities
# ============================================================================
@doc """
Check if a value is a Python object reference.
## Examples
{:ok, path} = SnakeBridge.call("pathlib", "Path", ["."])
SnakeBridge.ref?(path)
# => true
SnakeBridge.ref?("string")
# => false
"""
@spec ref?(term()) :: boolean()
defdelegate ref?(value), to: Ref
# ============================================================================
# Helpers & Macros (Existing)
# ============================================================================
@doc """
Call a helper function.
"""
defdelegate call_helper(helper, args \\ [], opts \\ []), to: Runtime
@doc """
Context manager macro for Python with statements.
"""
defmacro with_python(ref, do: block) do
quote do
require SnakeBridge.WithContext
SnakeBridge.WithContext.with_python(unquote(ref), do: unquote(block))
end
end
@doc """
Returns the SnakeBridge version.
"""
@spec version() :: String.t()
def version do
Application.spec(:snakebridge, :vsn) |> to_string()
end
end