tywrap 0.5.1 → 0.6.1

package/src/core/annotation-parser.ts

@@ -118,30 +118,35 @@ export function parseAnnotationToPythonType(
     return { kind: 'custom', name: n };
   };
 
+  // Extracts the leading quoted string from the body of a `Name(...)` factory call
+  // (the `inner` already stripped of the outer `Name(` and trailing `)`). Returns
+  // the unquoted name, or null when the body is not a well-formed quoted argument.
+  const extractQuotedFactoryArg = (inner: string): string | null => {
+    if (inner.length < 2) {
+      return null;
+    }
+
+    const quote = inner[0];
+    if ((quote !== "'" && quote !== '"') || inner[inner.length - 1] !== quote) {
+      return null;
+    }
+
+    const commaIndex = inner.indexOf(',');
+    const quoted = commaIndex === -1 ? inner : inner.slice(0, commaIndex).trimEnd();
+    if (quoted.length < 2 || quoted[quoted.length - 1] !== quote) {
+      return null;
+    }
+
+    return quoted.slice(1, -1);
+  };
+
   const parseTypingFactoryName = (text: string, name: string): string | null => {
     for (const prefix of modulePrefixes) {
       const start = `${prefix}${name}(`;
       if (!text.startsWith(start) || !text.endsWith(')')) {
         continue;
       }
-
-      const inner = text.slice(start.length, -1).trim();
-      if (inner.length < 2) {
-        return null;
-      }
-
-      const quote = inner[0];
-      if ((quote !== "'" && quote !== '"') || inner[inner.length - 1] !== quote) {
-        return null;
-      }
-
-      const commaIndex = inner.indexOf(',');
-      const quoted = commaIndex === -1 ? inner : inner.slice(0, commaIndex).trimEnd();
-      if (quoted.length < 2 || quoted[quoted.length - 1] !== quote) {
-        return null;
-      }
-
-      return quoted.slice(1, -1);
+      return extractQuotedFactoryArg(text.slice(start.length, -1).trim());
     }
 
     return null;
@@ -245,6 +250,19 @@ export function parseAnnotationToPythonType(
     return results;
   };
 
+  // Returns the content between the first `[` and the last `]` of a special-form
+  // annotation (e.g. the `int, str` of `Union[int, str]`). Callers gate this on a
+  // prefix check, so the brackets are known to be present.
+  const bracketInner = (raw: string): string => raw.slice(raw.indexOf('[') + 1, raw.lastIndexOf(']'));
+
+  // True when `raw` opens with one of the supported module prefixes followed by
+  // `name[`. Mirrors the inlined `raw.startsWith('typing.Name[') || ...` checks so
+  // the branch ordering in `parse` is unchanged.
+  const startsWithSpecialForm = (raw: string, name: string): boolean =>
+    raw.startsWith(`typing.${name}[`) ||
+    raw.startsWith(`typing_extensions.${name}[`) ||
+    raw.startsWith(`${name}[`);
+
   const splitGenericInvocation = (raw: string): { name: string; inner: string } | null => {
     if (!raw.endsWith(']')) {
       return null;
@@ -277,6 +295,86 @@ export function parseAnnotationToPythonType(
     };
   };
 
+  // Parses the body of a `Callable[...]` form (the `inner` already stripped of the
+  // outer brackets). `recurse` is the top-level `parse`, threaded in so this stays a
+  // standalone helper. Returns null only when the form has fewer than the required
+  // two comma-separated sections, matching the original guarded branch exactly.
+  const parseCallableInner = (
+    inner: string,
+    depth: number,
+    recurse: (ann: unknown, depth?: number) => PythonType
+  ): PythonType | null => {
+    const parts = splitTopLevel(inner, ',');
+    if (parts.length < 2) {
+      return null;
+    }
+
+    const paramsPart = (parts[0] ?? '').trim();
+    const returnPart = parts.slice(1).join(',').trim();
+    const returnType = recurse(returnPart, depth + 1);
+
+    if (paramsPart === '...' || paramsPart === 'Ellipsis') {
+      return {
+        kind: 'callable',
+        parameters: [{ kind: 'custom', name: '...' }],
+        returnType,
+      };
+    }
+
+    const directParamSpec = recurse(paramsPart, depth + 1);
+    if (directParamSpec.kind === 'paramspec') {
+      return {
+        kind: 'callable',
+        parameters: [],
+        parameterSpec: directParamSpec,
+        returnType,
+      };
+    }
+
+    const paramInner =
+      paramsPart.startsWith('[') && paramsPart.endsWith(']') ? paramsPart.slice(1, -1) : '';
+    const trimmed = paramInner.trim();
+    if (trimmed === '...' || trimmed === 'Ellipsis') {
+      return {
+        kind: 'callable',
+        parameters: [{ kind: 'custom', name: '...' }],
+        returnType,
+      };
+    }
+
+    const parameters = trimmed
+      ? splitTopLevel(trimmed, ',').map(p => recurse(p.trim(), depth + 1))
+      : [];
+    return { kind: 'callable', parameters, returnType };
+  };
+
+  // Handles the `typing.TypeVar('T')` / `ParamSpec('P')` / `TypeVarTuple('Ts')`
+  // factory-call forms. Returns null when `rawText` is none of them, preserving
+  // the original sequential fall-through (TypeVar, then ParamSpec, then TypeVarTuple).
+  const parseTypingFactory = (rawText: string): PythonType | null => {
+    const typeVarName = parseTypingFactoryName(rawText, 'TypeVar');
+    if (typeVarName) {
+      return mapKnownTypeParameter(typeVarName) ?? { kind: 'typevar', name: typeVarName };
+    }
+
+    const paramSpecName = parseTypingFactoryName(rawText, 'ParamSpec');
+    if (paramSpecName) {
+      return mapKnownTypeParameter(paramSpecName) ?? { kind: 'paramspec', name: paramSpecName };
+    }
+
+    const typeVarTupleName = parseTypingFactoryName(rawText, 'TypeVarTuple');
+    if (typeVarTupleName) {
+      return (
+        mapKnownTypeParameter(typeVarTupleName) ?? {
+          kind: 'typevartuple',
+          name: typeVarTupleName,
+        }
+      );
+    }
+
+    return null;
+  };
+
   const parse = (ann: unknown, depth = 0): PythonType => {
     if (ann === null || ann === undefined) {
       return unknownType();
@@ -320,33 +418,18 @@ export function parseAnnotationToPythonType(
       }
     }
 
-    if (
-      raw.startsWith('typing.Union[') ||
-      raw.startsWith('typing_extensions.Union[') ||
-      raw.startsWith('Union[')
-    ) {
-      const inner = raw.slice(raw.indexOf('[') + 1, raw.lastIndexOf(']'));
-      const parts = splitTopLevel(inner, ',');
+    if (startsWithSpecialForm(raw, 'Union')) {
+      const parts = splitTopLevel(bracketInner(raw), ',');
       const types = parts.map(p => parse(p.trim(), depth + 1));
       return { kind: 'union', types };
     }
 
-    if (
-      raw.startsWith('typing.Optional[') ||
-      raw.startsWith('typing_extensions.Optional[') ||
-      raw.startsWith('Optional[')
-    ) {
-      const inner = raw.slice(raw.indexOf('[') + 1, raw.lastIndexOf(']'));
-      return { kind: 'optional', type: parse(inner, depth + 1) };
+    if (startsWithSpecialForm(raw, 'Optional')) {
+      return { kind: 'optional', type: parse(bracketInner(raw), depth + 1) };
     }
 
-    if (
-      raw.startsWith('typing.Literal[') ||
-      raw.startsWith('typing_extensions.Literal[') ||
-      raw.startsWith('Literal[')
-    ) {
-      const inner = raw.slice(raw.indexOf('[') + 1, raw.lastIndexOf(']'));
-      const parts = splitTopLevel(inner, ',');
+    if (startsWithSpecialForm(raw, 'Literal')) {
+      const parts = splitTopLevel(bracketInner(raw), ',');
       if (parts.length === 1) {
         return mapLiteral(String(parts[0] ?? '').trim());
       }
@@ -360,24 +443,9 @@ export function parseAnnotationToPythonType(
       return parse(extMatch[3] ?? '', depth + 1);
     }
 
-    const typeVarName = parseTypingFactoryName(rawText, 'TypeVar');
-    if (typeVarName) {
-      return mapKnownTypeParameter(typeVarName) ?? { kind: 'typevar', name: typeVarName };
-    }
-
-    const paramSpecName = parseTypingFactoryName(rawText, 'ParamSpec');
-    if (paramSpecName) {
-      return mapKnownTypeParameter(paramSpecName) ?? { kind: 'paramspec', name: paramSpecName };
-    }
-
-    const typeVarTupleName = parseTypingFactoryName(rawText, 'TypeVarTuple');
-    if (typeVarTupleName) {
-      return (
-        mapKnownTypeParameter(typeVarTupleName) ?? {
-          kind: 'typevartuple',
-          name: typeVarTupleName,
-        }
-      );
+    const typingFactory = parseTypingFactory(rawText);
+    if (typingFactory) {
+      return typingFactory;
     }
     if (
       raw === 'typing.LiteralString' ||
@@ -387,61 +455,15 @@ export function parseAnnotationToPythonType(
       return { kind: 'primitive', name: 'str' };
     }
 
-    if (
-      raw.startsWith('typing.Callable[') ||
-      raw.startsWith('typing_extensions.Callable[') ||
-      raw.startsWith('Callable[')
-    ) {
-      const inner = raw.slice(raw.indexOf('[') + 1, raw.lastIndexOf(']'));
-      const parts = splitTopLevel(inner, ',');
-      if (parts.length >= 2) {
-        const paramsPart = (parts[0] ?? '').trim();
-        const returnPart = parts.slice(1).join(',').trim();
-        const returnType = parse(returnPart, depth + 1);
-
-        if (paramsPart === '...' || paramsPart === 'Ellipsis') {
-          return {
-            kind: 'callable',
-            parameters: [{ kind: 'custom', name: '...' }],
-            returnType,
-          };
-        }
-
-        const directParamSpec = parse(paramsPart, depth + 1);
-        if (directParamSpec.kind === 'paramspec') {
-          return {
-            kind: 'callable',
-            parameters: [],
-            parameterSpec: directParamSpec,
-            returnType,
-          };
-        }
-
-        const paramInner =
-          paramsPart.startsWith('[') && paramsPart.endsWith(']') ? paramsPart.slice(1, -1) : '';
-        const trimmed = paramInner.trim();
-        if (trimmed === '...' || trimmed === 'Ellipsis') {
-          return {
-            kind: 'callable',
-            parameters: [{ kind: 'custom', name: '...' }],
-            returnType,
-          };
-        }
-
-        const parameters = trimmed
-          ? splitTopLevel(trimmed, ',').map(p => parse(p.trim(), depth + 1))
-          : [];
-        return { kind: 'callable', parameters, returnType };
+    if (startsWithSpecialForm(raw, 'Callable')) {
+      const callable = parseCallableInner(bracketInner(raw), depth, parse);
+      if (callable) {
+        return callable;
       }
     }
 
-    if (
-      raw.startsWith('typing.Mapping[') ||
-      raw.startsWith('typing_extensions.Mapping[') ||
-      raw.startsWith('Mapping[')
-    ) {
-      const inner = raw.slice(raw.indexOf('[') + 1, raw.lastIndexOf(']'));
-      const parts = splitTopLevel(inner, ',');
+    if (startsWithSpecialForm(raw, 'Mapping')) {
+      const parts = splitTopLevel(bracketInner(raw), ',');
       return {
         kind: 'collection',
         name: 'dict',
@@ -449,13 +471,8 @@ export function parseAnnotationToPythonType(
       };
     }
 
-    if (
-      raw.startsWith('typing.Annotated[') ||
-      raw.startsWith('typing_extensions.Annotated[') ||
-      raw.startsWith('Annotated[')
-    ) {
-      const inner = raw.slice(raw.indexOf('[') + 1, raw.lastIndexOf(']'));
-      const parts = splitTopLevel(inner, ',');
+    if (startsWithSpecialForm(raw, 'Annotated')) {
+      const parts = splitTopLevel(bracketInner(raw), ',');
       if (parts.length > 0) {
         return {
           kind: 'annotated',
@@ -465,13 +482,8 @@ export function parseAnnotationToPythonType(
       }
     }
 
-    if (
-      raw.startsWith('typing.Unpack[') ||
-      raw.startsWith('typing_extensions.Unpack[') ||
-      raw.startsWith('Unpack[')
-    ) {
-      const inner = raw.slice(raw.indexOf('[') + 1, raw.lastIndexOf(']'));
-      return { kind: 'unpack', type: parse(inner, depth + 1) };
+    if (startsWithSpecialForm(raw, 'Unpack')) {
+      return { kind: 'unpack', type: parse(bracketInner(raw), depth + 1) };
     }
 
     const coll = normalizeCollectionName(raw);

package/src/core/discovery.ts

@@ -12,6 +12,32 @@ import { getComponentLogger } from '../utils/logger.js';
 
 const log = getComponentLogger('Discovery');
 
+/**
+ * Strict allow-list for Python module names.
+ *
+ * Module names are interpolated into a `python -c` source string, so anything
+ * outside a conservative identifier-plus-dots shape (e.g. quotes, semicolons,
+ * newlines, parentheses, whitespace) could break out of the intended
+ * expression and inject arbitrary code into the interpreter. Reject those
+ * before they ever reach the subprocess argv.
+ */
+const SAFE_MODULE_NAME = /^[A-Za-z_][A-Za-z0-9_.]*$/;
+
+/**
+ * Validate a Python module name against {@link SAFE_MODULE_NAME}.
+ *
+ * @throws {Error} if the name is empty or contains characters that could
+ *   escape the `python -c` expression it is interpolated into.
+ */
+function assertSafeModuleName(moduleName: string): void {
+  if (typeof moduleName !== 'string' || !SAFE_MODULE_NAME.test(moduleName)) {
+    throw new Error(
+      `Unsafe Python module name rejected: ${JSON.stringify(moduleName)}. ` +
+        'Module names must match /^[A-Za-z_][A-Za-z0-9_.]*$/.'
+    );
+  }
+}
+
 export interface ModuleInfo {
   name: string;
   path: string;
@@ -87,44 +113,70 @@ export class ModuleDiscovery {
    * Resolve Python module path from module name
    */
   async resolvePythonPath(moduleName: string): Promise<string | null> {
+    // Reject names that could break out of the `python -c` expression below.
+    assertSafeModuleName(moduleName);
+
     // Check cache first
     const cached = this.moduleCache.get(moduleName);
     if (cached) {
       return cached.path;
     }
 
-    // Try to resolve using Python's module finder
-    if (processUtils.isAvailable()) {
-      try {
-        const pythonPath = await this.getPythonExecutable();
-        const result = await processUtils.exec(
-          pythonPath,
-          [
-            '-c',
-            `import ${moduleName}; print(${moduleName}.__file__ if hasattr(${moduleName}, '__file__') else '${moduleName}.__path__[0]' if hasattr(${moduleName}, '__path__') else 'builtin')`,
-          ],
-          { timeoutMs: this.options.timeoutMs }
-        );
+    // Try to resolve using Python's module finder, then fall back to a
+    // filesystem search across the known Python paths.
+    return (
+      (await this.resolveViaInterpreter(moduleName)) ??
+      (await this.resolveViaSearchPaths(moduleName))
+    );
+  }
+
+  /**
+   * Resolve a module path by asking the Python interpreter for `__file__`.
+   * Caches and returns the resolved path, or `null` when the interpreter is
+   * unavailable, reports a builtin, or fails.
+   */
+  private async resolveViaInterpreter(moduleName: string): Promise<string | null> {
+    if (!processUtils.isAvailable()) {
+      return null;
+    }
+
+    try {
+      const pythonPath = await this.getPythonExecutable();
+      const result = await processUtils.exec(
+        pythonPath,
+        [
+          '-c',
+          `import ${moduleName}; print(${moduleName}.__file__ if hasattr(${moduleName}, '__file__') else '${moduleName}.__path__[0]' if hasattr(${moduleName}, '__path__') else 'builtin')`,
+        ],
+        { timeoutMs: this.options.timeoutMs }
+      );
 
-        if (result.code === 0 && result.stdout.trim() !== 'builtin') {
-          const modulePath = result.stdout.trim();
+      if (result.code === 0 && result.stdout.trim() !== 'builtin') {
+        const modulePath = result.stdout.trim();
 
-          // Cache the result
-          this.moduleCache.set(moduleName, {
-            name: moduleName,
-            path: modulePath,
-            isPackage: modulePath.endsWith('__init__.py'),
-            dependencies: [],
-          });
+        // Cache the result
+        this.moduleCache.set(moduleName, {
+          name: moduleName,
+          path: modulePath,
+          isPackage: modulePath.endsWith('__init__.py'),
+          dependencies: [],
+        });
 
-          return modulePath;
-        }
-      } catch (error) {
-        log.warn('Failed to resolve module', { moduleName, error: String(error) });
+        return modulePath;
       }
+    } catch (error) {
+      log.warn('Failed to resolve module', { moduleName, error: String(error) });
     }
 
-    // Fallback: search in known Python paths
+    return null;
+  }
+
+  /**
+   * Fallback resolution: probe `<path>/<module>.py` and
+   * `<path>/<module>/__init__.py` across the known Python search paths.
+   * Returns the first existing candidate, or `null`.
+   */
+  private async resolveViaSearchPaths(moduleName: string): Promise<string | null> {
     const searchPaths = await this.getPythonSearchPaths();
 
     for (const searchPath of searchPaths) {
@@ -208,40 +260,33 @@ export class ModuleDiscovery {
    */
   parseDependenciesFromSource(source: string): string[] {
     const dependencies: string[] = [];
-    const lines = source.split('\n');
 
-    for (const line of lines) {
-      const trimmed = line.trim();
-
-      // Skip comments and empty lines
-      if (trimmed.startsWith('#') || trimmed === '') {
-        continue;
+    for (const line of source.split('\n')) {
+      const moduleName = this.extractImportedModule(line.trim());
+      if (moduleName) {
+        dependencies.push(moduleName);
       }
+    }
 
-      // Match import statements
-      const importMatch = trimmed.match(/^import\s+([a-zA-Z_][a-zA-Z0-9_\.]*)/);
-      if (importMatch?.[1]) {
-        const fullModuleName = importMatch[1];
-        const moduleName = fullModuleName.split('.')[0]; // Get top-level module
-        if (moduleName) {
-          dependencies.push(moduleName);
-        }
-        continue;
-      }
+    return [...new Set(dependencies)]; // Remove duplicates
+  }
 
-      // Match from ... import statements
-      const fromMatch = trimmed.match(/^from\s+([a-zA-Z_][a-zA-Z0-9_\.]*)\s+import/);
-      if (fromMatch?.[1]) {
-        const fullModuleName = fromMatch[1];
-        const moduleName = fullModuleName.split('.')[0]; // Get top-level module
-        if (moduleName) {
-          dependencies.push(moduleName);
-        }
-        continue;
-      }
+  /**
+   * Extract the top-level module name imported by a single (trimmed) source
+   * line, or `null` when the line is a comment, blank, or not an import.
+   */
+  private extractImportedModule(trimmed: string): string | null {
+    // Skip comments and empty lines
+    if (trimmed.startsWith('#') || trimmed === '') {
+      return null;
     }
 
-    return [...new Set(dependencies)]; // Remove duplicates
+    // Match `import x.y.z` and `from x.y.z import ...` statements.
+    const match =
+      trimmed.match(/^import\s+([a-zA-Z_][a-zA-Z0-9_\.]*)/) ??
+      trimmed.match(/^from\s+([a-zA-Z_][a-zA-Z0-9_\.]*)\s+import/);
+
+    return match?.[1]?.split('.')[0] ?? null; // Get top-level module
   }
 
   /**
@@ -383,6 +428,9 @@ export class ModuleDiscovery {
    * Get module version information
    */
   async getModuleVersion(moduleName: string): Promise<string | undefined> {
+    // Reject names that could break out of the `python -c` expression below.
+    assertSafeModuleName(moduleName);
+
     if (!processUtils.isAvailable()) {
       return undefined;
     }