smartregistry/ts/pypi/classes.pypiupstream.ts

import * as plugins from '../plugins.js';
import { BaseUpstream } from '../upstream/classes.baseupstream.js';
import type {
  IProtocolUpstreamConfig,
  IUpstreamFetchContext,
  IUpstreamRegistryConfig,
} from '../upstream/interfaces.upstream.js';

/**
 * PyPI-specific upstream implementation.
 *
 * Handles:
 * - Simple API (HTML) - PEP 503
 * - JSON API - PEP 691
 * - Package file downloads (wheels, sdists)
 * - Package name normalization
 */
export class PypiUpstream extends BaseUpstream {
  protected readonly protocolName = 'pypi';

  /** Local registry URL for rewriting download URLs */
  private readonly localRegistryUrl: string;

  constructor(
    config: IProtocolUpstreamConfig,
    localRegistryUrl: string,
    logger?: plugins.smartlog.Smartlog,
  ) {
    super(config, logger);
    this.localRegistryUrl = localRegistryUrl;
  }

  /**
   * Fetch Simple API index (list of all packages) in HTML format.
   */
  public async fetchSimpleIndex(): Promise<string | null> {
    const context: IUpstreamFetchContext = {
      protocol: 'pypi',
      resource: '*',
      resourceType: 'index',
      path: '/simple/',
      method: 'GET',
      headers: {
        'accept': 'text/html',
      },
      query: {},
    };

    const result = await this.fetch(context);

    if (!result || !result.success) {
      return null;
    }

    if (Buffer.isBuffer(result.body)) {
      return result.body.toString('utf8');
    }

    return typeof result.body === 'string' ? result.body : null;
  }

  /**
   * Fetch Simple API package page (list of files) in HTML format.
   */
  public async fetchSimplePackage(packageName: string): Promise<string | null> {
    const normalizedName = this.normalizePackageName(packageName);
    const path = `/simple/${normalizedName}/`;

    const context: IUpstreamFetchContext = {
      protocol: 'pypi',
      resource: packageName,
      resourceType: 'simple',
      path,
      method: 'GET',
      headers: {
        'accept': 'text/html',
      },
      query: {},
    };

    const result = await this.fetch(context);

    if (!result || !result.success) {
      return null;
    }

    if (Buffer.isBuffer(result.body)) {
      return result.body.toString('utf8');
    }

    return typeof result.body === 'string' ? result.body : null;
  }

  /**
   * Fetch package metadata using JSON API (PEP 691).
   */
  public async fetchPackageJson(packageName: string): Promise<any | null> {
    const normalizedName = this.normalizePackageName(packageName);
    const path = `/simple/${normalizedName}/`;

    const context: IUpstreamFetchContext = {
      protocol: 'pypi',
      resource: packageName,
      resourceType: 'metadata',
      path,
      method: 'GET',
      headers: {
        'accept': 'application/vnd.pypi.simple.v1+json',
      },
      query: {},
    };

    const result = await this.fetch(context);

    if (!result || !result.success) {
      return null;
    }

    if (Buffer.isBuffer(result.body)) {
      return JSON.parse(result.body.toString('utf8'));
    }

    return result.body;
  }

  /**
   * Fetch full package info from PyPI JSON API (/pypi/{package}/json).
   */
  public async fetchPypiJson(packageName: string): Promise<any | null> {
    const normalizedName = this.normalizePackageName(packageName);
    const path = `/pypi/${normalizedName}/json`;

    const context: IUpstreamFetchContext = {
      protocol: 'pypi',
      resource: packageName,
      resourceType: 'pypi-json',
      path,
      method: 'GET',
      headers: {
        'accept': 'application/json',
      },
      query: {},
    };

    const result = await this.fetch(context);

    if (!result || !result.success) {
      return null;
    }

    if (Buffer.isBuffer(result.body)) {
      return JSON.parse(result.body.toString('utf8'));
    }

    return result.body;
  }

  /**
   * Fetch a package file (wheel or sdist) from upstream.
   */
  public async fetchPackageFile(packageName: string, filename: string): Promise<Buffer | null> {
    const normalizedName = this.normalizePackageName(packageName);
    const path = `/packages/${normalizedName}/${filename}`;

    const context: IUpstreamFetchContext = {
      protocol: 'pypi',
      resource: packageName,
      resourceType: 'package',
      path,
      method: 'GET',
      headers: {
        'accept': 'application/octet-stream',
      },
      query: {},
    };

    const result = await this.fetch(context);

    if (!result || !result.success) {
      return null;
    }

    return Buffer.isBuffer(result.body) ? result.body : Buffer.from(result.body);
  }

  /**
   * Normalize a PyPI package name according to PEP 503.
   * - Lowercase all characters
   * - Replace runs of ., -, _ with single -
   */
  private normalizePackageName(name: string): string {
    return name.toLowerCase().replace(/[-_.]+/g, '-');
  }

  /**
   * Override URL building for PyPI-specific handling.
   */
  protected buildUpstreamUrl(
    upstream: IUpstreamRegistryConfig,
    context: IUpstreamFetchContext,
  ): string {
    let baseUrl = upstream.url;

    // Remove trailing slash
    if (baseUrl.endsWith('/')) {
      baseUrl = baseUrl.slice(0, -1);
    }

    return `${baseUrl}${context.path}`;
  }
}