@mui/internal-docs-infra/pipeline/loadIsomorphicCodeVariant/enhanceCode.mjs

MUI Infra - internal documentation creation tools.

/**
 * Helper to check if a source is a HAST root (already parsed)
 */
function isHastRoot(source) {
  return typeof source === 'object' && source !== null && 'type' in source && source.type === 'root';
}

/**
 * Async function to enhance parsed code variants and their extraFiles.
 * Applies sourceEnhancers to HAST nodes, using comments stored in the variant.
 */
export async function enhanceCode(code, sourceEnhancers) {
  if (!sourceEnhancers || sourceEnhancers.length === 0) {
    return code;
  }

  /**
   * Helper to apply enhancers sequentially to a HAST root
   */
  async function applyEnhancers(source, comments, fileName) {
    return sourceEnhancers.reduce(async (accPromise, enhancer) => {
      const acc = await accPromise;
      return enhancer(acc, comments, fileName);
    }, Promise.resolve(source));
  }

  /**
   * Helper to enhance a single variant
   */
  async function enhanceVariant(variantCode) {
    if (typeof variantCode === 'string') {
      return variantCode;
    }
    if (!variantCode.source || !isHastRoot(variantCode.source)) {
      return variantCode;
    }

    // Apply enhancers to the main source
    const fileName = variantCode.fileName || 'unknown';
    const enhancedSource = await applyEnhancers(variantCode.source, variantCode.comments, fileName);

    // Also enhance extraFiles if they have HAST sources
    let enhancedExtraFiles = variantCode.extraFiles;
    if (variantCode.extraFiles) {
      const extraFileEntries = await Promise.all(Object.entries(variantCode.extraFiles).map(async ([extraFileName, fileContent]) => {
        if (typeof fileContent === 'string') {
          return [extraFileName, fileContent]; // Keep string as-is
        }
        if (fileContent && typeof fileContent === 'object' && isHastRoot(fileContent.source)) {
          // Apply enhancers to this extra file's source
          const enhancedExtraSource = await applyEnhancers(fileContent.source, fileContent.comments, extraFileName);
          return [extraFileName, {
            ...fileContent,
            source: enhancedExtraSource,
            // Clear comments after enhancing since they've been consumed
            comments: undefined
          }];
        }
        return [extraFileName, fileContent]; // Keep as-is for other cases
      }));
      enhancedExtraFiles = Object.fromEntries(extraFileEntries);
    }
    return {
      ...variantCode,
      source: enhancedSource,
      extraFiles: enhancedExtraFiles,
      // Clear comments after enhancing since they've been consumed
      comments: undefined
    };
  }

  // Process all variants in parallel
  const entries = Object.entries(code);
  const enhancedEntries = await Promise.all(entries.map(async ([variant, variantCode]) => {
    if (!variantCode) {
      return [variant, variantCode];
    }
    const enhanced = await enhanceVariant(variantCode);
    return [variant, enhanced];
  }));
  return Object.fromEntries(enhancedEntries);
}