283 lines
11 KiB
Zig
283 lines
11 KiB
Zig
const std = @import("std");
|
|
const zfin = @import("../root.zig");
|
|
const cli = @import("common.zig");
|
|
const framework = @import("framework.zig");
|
|
const fmt = @import("../format.zig");
|
|
const isCusipLike = @import("../models/portfolio.zig").isCusipLike;
|
|
|
|
pub const ParsedArgs = struct {
|
|
/// Either a symbol (e.g. "AAPL") or a path to a portfolio file
|
|
/// (e.g. "portfolio.srf"). Distinguished by suffix / path-separator
|
|
/// heuristic at run time so the user can pass either form.
|
|
arg: []const u8,
|
|
};
|
|
|
|
pub const meta: framework.Meta = .{
|
|
.name = "enrich",
|
|
.group = .hygiene,
|
|
.synopsis = "Bootstrap metadata.srf from Alpha Vantage (25 req/day limit)",
|
|
.help =
|
|
\\Usage: zfin enrich <FILE|SYMBOL>
|
|
\\
|
|
\\Bootstrap a `metadata.srf` classification file from Alpha
|
|
\\Vantage's OVERVIEW endpoint. Two modes:
|
|
\\
|
|
\\ - File mode (path or `*.srf` suffix): enrich every stock
|
|
\\ symbol in the portfolio. Output is a complete SRF file
|
|
\\ written to stdout — redirect into metadata.srf and
|
|
\\ edit by hand for accuracy.
|
|
\\ - Symbol mode (anything else): enrich a single symbol and
|
|
\\ emit one appendable SRF line. Useful for adding to an
|
|
\\ existing metadata.srf without rerunning the whole file.
|
|
\\
|
|
\\Caveats: Alpha Vantage's free tier is 25 requests/day. The
|
|
\\OVERVIEW data is US-domicile-biased — international ETFs
|
|
\\classify as `geo::US`. Always review the output before
|
|
\\saving as `metadata.srf`. Requires ALPHAVANTAGE_API_KEY.
|
|
\\
|
|
\\Examples:
|
|
\\ zfin enrich portfolio.srf > metadata.srf # whole portfolio
|
|
\\ zfin enrich AAPL >> metadata.srf # single symbol append
|
|
\\
|
|
,
|
|
.uppercase_first_arg = false,
|
|
};
|
|
|
|
comptime {
|
|
framework.validateCommandModule(@This());
|
|
}
|
|
|
|
pub fn parseArgs(ctx: *framework.RunCtx, cmd_args: []const []const u8) !ParsedArgs {
|
|
if (cmd_args.len < 1) {
|
|
try cli.stderrPrint(ctx.io, "Error: 'enrich' requires a portfolio file path or symbol\n");
|
|
return error.MissingArg;
|
|
}
|
|
if (cmd_args.len > 1) {
|
|
try cli.stderrPrint(ctx.io, "Error: 'enrich' takes a single argument (file path or symbol)\n");
|
|
return error.UnexpectedArg;
|
|
}
|
|
return .{ .arg = cmd_args[0] };
|
|
}
|
|
|
|
const OverviewMeta = struct {
|
|
sector: []const u8,
|
|
geo: []const u8,
|
|
asset_class: []const u8,
|
|
};
|
|
|
|
/// Derive sector, geo, and asset_class from an Alpha Vantage company overview.
|
|
fn deriveMetadata(overview: zfin.CompanyOverview, sector_buf: []u8) OverviewMeta {
|
|
const sector_raw = overview.sector orelse "Unknown";
|
|
const sector_str = cli.fmt.toTitleCase(sector_buf, sector_raw);
|
|
const country_str = overview.country orelse "US";
|
|
const geo_str = if (std.mem.eql(u8, country_str, "USA")) "US" else country_str;
|
|
|
|
const asset_class_str = blk: {
|
|
if (overview.asset_type) |at| {
|
|
if (std.mem.eql(u8, at, "ETF")) break :blk "ETF";
|
|
if (std.mem.eql(u8, at, "Mutual Fund")) break :blk "Mutual Fund";
|
|
}
|
|
if (overview.market_cap) |mc_str| {
|
|
const mc = std.fmt.parseInt(u64, mc_str, 10) catch 0;
|
|
if (mc >= 10_000_000_000) break :blk "US Large Cap";
|
|
if (mc >= 2_000_000_000) break :blk "US Mid Cap";
|
|
break :blk "US Small Cap";
|
|
}
|
|
break :blk "US Large Cap";
|
|
};
|
|
|
|
return .{ .sector = sector_str, .geo = geo_str, .asset_class = asset_class_str };
|
|
}
|
|
|
|
/// CLI `enrich` command: bootstrap a metadata.srf file from Alpha Vantage OVERVIEW data.
|
|
/// Reads the portfolio, extracts stock symbols, fetches sector/industry/country for each,
|
|
/// and outputs a metadata SRF file to stdout.
|
|
/// If the argument looks like a symbol (no path separators, no .srf extension), enrich just that symbol.
|
|
pub fn run(ctx: *framework.RunCtx, parsed: ParsedArgs) !void {
|
|
const svc = ctx.svc orelse return error.MissingDataService;
|
|
// Determine if arg is a symbol or a file path
|
|
const is_file = std.mem.endsWith(u8, parsed.arg, ".srf") or
|
|
std.mem.indexOfScalar(u8, parsed.arg, '/') != null or
|
|
std.mem.indexOfScalar(u8, parsed.arg, '.') != null;
|
|
|
|
if (!is_file) {
|
|
// Single symbol mode: enrich one symbol, output appendable SRF (no header)
|
|
try enrichSymbol(ctx.io, ctx.allocator, svc, parsed.arg, ctx.out);
|
|
return;
|
|
}
|
|
|
|
// Portfolio file mode: enrich all symbols
|
|
try enrichPortfolio(ctx.io, ctx.allocator, svc, parsed.arg, ctx.today, ctx.out);
|
|
}
|
|
|
|
/// Enrich a single symbol and output appendable SRF lines to stdout.
|
|
fn enrichSymbol(io: std.Io, allocator: std.mem.Allocator, svc: *zfin.DataService, sym: []const u8, out: *std.Io.Writer) !void {
|
|
{
|
|
var msg_buf: [128]u8 = undefined;
|
|
const msg = std.fmt.bufPrint(&msg_buf, " Fetching {s}...\n", .{sym}) catch " ...\n";
|
|
try cli.stderrPrint(io, msg);
|
|
}
|
|
|
|
const overview = svc.getCompanyOverview(sym) catch |err| {
|
|
if (err == zfin.DataError.NoApiKey) {
|
|
try cli.stderrPrint(io, "Error: ALPHAVANTAGE_API_KEY not set. Add it to .env\n");
|
|
return;
|
|
}
|
|
try cli.stderrPrint(io, "Error: Failed to fetch data for symbol\n");
|
|
try out.print("# {s} -- fetch failed\n", .{sym});
|
|
try out.print("# symbol::{s},sector::TODO,geo::TODO,asset_class::TODO\n", .{sym});
|
|
return;
|
|
};
|
|
defer {
|
|
if (overview.name) |n| allocator.free(n);
|
|
if (overview.sector) |s| allocator.free(s);
|
|
if (overview.industry) |ind| allocator.free(ind);
|
|
if (overview.country) |c| allocator.free(c);
|
|
if (overview.market_cap) |mc| allocator.free(mc);
|
|
if (overview.asset_type) |at| allocator.free(at);
|
|
}
|
|
|
|
var sector_buf: [64]u8 = undefined;
|
|
const derived = deriveMetadata(overview, §or_buf);
|
|
|
|
if (overview.name) |name| {
|
|
try out.print("# {s}\n", .{name});
|
|
}
|
|
try out.print("symbol::{s},sector::{s},geo::{s},asset_class::{s}\n", .{
|
|
sym, derived.sector, derived.geo, derived.asset_class,
|
|
});
|
|
}
|
|
|
|
/// Enrich all symbols from a portfolio file.
|
|
fn enrichPortfolio(io: std.Io, allocator: std.mem.Allocator, svc: *zfin.DataService, file_path: []const u8, as_of: zfin.Date, out: *std.Io.Writer) !void {
|
|
|
|
// Load portfolio
|
|
const file_data = std.Io.Dir.cwd().readFileAlloc(io, file_path, allocator, .limited(10 * 1024 * 1024)) catch {
|
|
try cli.stderrPrint(io, "Error: Cannot read portfolio file\n");
|
|
return;
|
|
};
|
|
defer allocator.free(file_data);
|
|
|
|
var portfolio = zfin.cache.deserializePortfolio(allocator, file_data) catch {
|
|
try cli.stderrPrint(io, "Error: Cannot parse portfolio file\n");
|
|
return;
|
|
};
|
|
defer portfolio.deinit();
|
|
|
|
// Get unique stock symbols (using display-oriented names)
|
|
const positions = try portfolio.positions(as_of, allocator);
|
|
defer allocator.free(positions);
|
|
|
|
// Get unique price symbols (raw API symbols)
|
|
const syms = try portfolio.stockSymbols(allocator);
|
|
defer allocator.free(syms);
|
|
|
|
try out.print("#!srfv1\n", .{});
|
|
try out.print("# Portfolio classification metadata\n", .{});
|
|
try out.print("# Generated from Alpha Vantage OVERVIEW data\n", .{});
|
|
try out.print("# Edit as needed: sector, geo, asset_class, pct:num:N\n", .{});
|
|
try out.print("#\n", .{});
|
|
try out.print("# For ETFs/funds with multi-class exposure, add multiple lines\n", .{});
|
|
try out.print("# with pct:num: values that sum to ~100\n\n", .{});
|
|
|
|
var success: usize = 0;
|
|
var skipped: usize = 0;
|
|
var failed: usize = 0;
|
|
|
|
for (syms, 0..) |sym, i| {
|
|
// Skip CUSIPs and known non-stock symbols
|
|
if (isCusipLike(sym)) {
|
|
// Find the display name for this CUSIP
|
|
const display: []const u8 = sym;
|
|
var note: ?[]const u8 = null;
|
|
for (positions) |pos| {
|
|
if (std.mem.eql(u8, pos.symbol, sym)) {
|
|
if (pos.note) |n| {
|
|
note = n;
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
try out.print("# CUSIP {s}", .{sym});
|
|
if (note) |n| try out.print(" ({s})", .{n});
|
|
try out.print(" -- fill in manually\n", .{});
|
|
try out.print("# symbol::{s},asset_class::TODO,geo::TODO\n\n", .{display});
|
|
skipped += 1;
|
|
continue;
|
|
}
|
|
|
|
// Progress to stderr
|
|
{
|
|
var msg_buf: [128]u8 = undefined;
|
|
const msg = std.fmt.bufPrint(&msg_buf, " [{d}/{d}] {s}...\n", .{ i + 1, syms.len, sym }) catch " ...\n";
|
|
try cli.stderrPrint(io, msg);
|
|
}
|
|
|
|
const overview = svc.getCompanyOverview(sym) catch {
|
|
try out.print("# {s} -- fetch failed\n", .{sym});
|
|
try out.print("# symbol::{s},sector::TODO,geo::TODO,asset_class::TODO\n\n", .{sym});
|
|
failed += 1;
|
|
continue;
|
|
};
|
|
// Free allocated strings from overview when done
|
|
defer {
|
|
if (overview.name) |n| allocator.free(n);
|
|
if (overview.sector) |s| allocator.free(s);
|
|
if (overview.industry) |ind| allocator.free(ind);
|
|
if (overview.country) |c| allocator.free(c);
|
|
if (overview.market_cap) |mc| allocator.free(mc);
|
|
if (overview.asset_type) |at| allocator.free(at);
|
|
}
|
|
|
|
var sector_buf: [64]u8 = undefined;
|
|
const derived = deriveMetadata(overview, §or_buf);
|
|
|
|
// Comment with the name for readability
|
|
if (overview.name) |name| {
|
|
try out.print("# {s}\n", .{name});
|
|
}
|
|
try out.print("symbol::{s},sector::{s},geo::{s},asset_class::{s}\n\n", .{
|
|
sym, derived.sector, derived.geo, derived.asset_class,
|
|
});
|
|
success += 1;
|
|
}
|
|
|
|
// Summary comment
|
|
try out.print("# ---\n", .{});
|
|
try out.print("# Enriched {d} symbols ({d} success, {d} skipped, {d} failed)\n", .{
|
|
syms.len, success, skipped, failed,
|
|
});
|
|
try out.print("# Review and edit this file, then save as metadata.srf\n", .{});
|
|
}
|
|
|
|
// ── Tests ────────────────────────────────────────────────────
|
|
|
|
test "parseArgs: accepts a symbol argument" {
|
|
var ctx: framework.RunCtx = undefined;
|
|
ctx.io = std.testing.io;
|
|
const args = [_][]const u8{"AAPL"};
|
|
const parsed = try parseArgs(&ctx, &args);
|
|
try std.testing.expectEqualStrings("AAPL", parsed.arg);
|
|
}
|
|
|
|
test "parseArgs: accepts a file path argument" {
|
|
var ctx: framework.RunCtx = undefined;
|
|
ctx.io = std.testing.io;
|
|
const args = [_][]const u8{"portfolio.srf"};
|
|
const parsed = try parseArgs(&ctx, &args);
|
|
try std.testing.expectEqualStrings("portfolio.srf", parsed.arg);
|
|
}
|
|
|
|
test "parseArgs: missing arg errors" {
|
|
var ctx: framework.RunCtx = undefined;
|
|
ctx.io = std.testing.io;
|
|
const args = [_][]const u8{};
|
|
try std.testing.expectError(error.MissingArg, parseArgs(&ctx, &args));
|
|
}
|
|
|
|
test "parseArgs: extra args error" {
|
|
var ctx: framework.RunCtx = undefined;
|
|
ctx.io = std.testing.io;
|
|
const args = [_][]const u8{ "AAPL", "extra" };
|
|
try std.testing.expectError(error.UnexpectedArg, parseArgs(&ctx, &args));
|
|
}
|