Files
BIPAuto/utils/discrete_material_plan/extractor.py
Misaka_Company da567a2679 feat(extractor): Add DiscreteMaterialPlanExtractor class for batch data extraction
- Create DiscreteMaterialPlanExtractor class orchestrating web operations and post-processing
- Implement extract_and_process() combining download + Excel conversion
- Implement post_process_downloads() using ExcelConverter for merge
- Add extract_from_file() convenience function for file-based ID input
- Add component test file verifying class structure and methods
- Follow existing patterns: verbose logging, explicit parameters, English comments
2026-03-27 13:48:07 +08:00

335 lines
11 KiB
Python

"""
High-level extractor component for discrete material plan data extraction.
Provides batch processing with progress reporting and error handling.
"""
import tempfile
import shutil
from pathlib import Path
from typing import List, Optional
from playwright.sync_api import sync_playwright
from utils.auth import login
from .excel_converter import ExcelConverter
from .extractor_core import (
navigate_to_discrete_material_page,
setup_query_interface,
execute_batch_download_workflow,
)
class DiscreteMaterialPlanExtractor:
"""
High-level extractor for discrete material plan data.
Handles session management, batch processing, and post-processing.
"""
def __init__(
self,
username: str,
password: str,
base_url: str,
headless: bool = True,
ignore_https_errors: bool = True,
verbose: bool = True,
download_dir: Optional[str] = None,
batch_size: int = 10,
):
"""
Initialize the extractor.
Args:
username: Login username
password: Login password
base_url: Base URL for the ERP system
headless: Whether to run browser in headless mode
ignore_https_errors: Whether to ignore HTTPS certificate errors
verbose: Whether to print detailed logs
download_dir: Directory to save downloaded files (default: temp dir)
batch_size: Number of order IDs per batch
"""
self.username = username
self.password = password
self.base_url = base_url
self.headless = headless
self.ignore_https_errors = ignore_https_errors
self.verbose = verbose
self.batch_size = batch_size
# Set up download directory
if download_dir:
self.download_dir = Path(download_dir)
self.download_dir.mkdir(parents=True, exist_ok=True)
else:
self.download_dir = Path(tempfile.mkdtemp())
self._is_temp_dir = True
self._print(f"Download directory: {self.download_dir}")
def _print(self, *args, **kwargs):
"""Print log message if verbose mode is enabled."""
if self.verbose:
print(*args, **kwargs)
def _cleanup_temp_files(self):
"""Remove temporary files created during extraction."""
if hasattr(self, "_is_temp_dir") and self._is_temp_dir:
try:
shutil.rmtree(self.download_dir)
self._print(f"Cleaned up temp directory: {self.download_dir}")
except Exception as e:
self._print(f"Warning: Failed to clean up temp directory: {e}")
def _get_login_url(self) -> str:
"""
Construct the complete login URL.
Returns:
str: Complete login page URL
"""
base = self.base_url.rstrip("/")
return f"{base}/yonbip/resources/uap/rbac/login/main/index.html"
def _chunk_order_ids(self, order_ids: List[str]) -> List[List[str]]:
"""
Split order IDs into batches.
Args:
order_ids: List of order IDs to process
Returns:
List of batches, where each batch is a list of order IDs
"""
chunks = []
for i in range(0, len(order_ids), self.batch_size):
chunks.append(order_ids[i : i + self.batch_size])
return chunks
def _download_batches(
self, order_ids: List[str]
) -> List[str]:
"""
Download data for all batches of order IDs.
Args:
order_ids: List of order IDs to download
Returns:
List of paths to downloaded files
"""
downloaded_files = []
chunks = self._chunk_order_ids(order_ids)
self._print(f"Processing {len(order_ids)} order IDs in {len(chunks)} batches")
with sync_playwright() as playwright:
# Login
self._print("Logging in...")
browser, context, page, main_frame = login(
playwright=playwright,
username=self.username,
password=self.password,
url=self._get_login_url(),
headless=self.headless,
ignore_https_errors=self.ignore_https_errors,
verbose=self.verbose,
)
try:
# Navigate to discrete material page
self._print("Navigating to discrete material plan page...")
work_frame, page1 = navigate_to_discrete_material_page(
main_frame, page
)
# Setup query interface
self._print("Setting up query interface...")
setup_query_interface(work_frame)
# Process each batch
for batch_index, batch in enumerate(chunks):
self._print(
f"Processing batch {batch_index + 1}/{len(chunks)} ({len(batch)} orders)"
)
file_path = execute_batch_download_workflow(
work_frame=work_frame,
page=page1,
order_ids=batch,
batch_index=batch_index,
download_dir=str(self.download_dir),
)
downloaded_files.append(file_path)
self._print(f"Downloaded: {file_path}")
finally:
# Cleanup
self._print("Closing browser...")
context.close()
browser.close()
return downloaded_files
def post_process_downloads(
self, downloaded_files: List[str], output_file: Optional[str] = None
) -> Optional[str]:
"""
Convert and merge downloaded Excel files into structured format.
Args:
downloaded_files: List of paths to downloaded Excel files
output_file: Path to save merged result (default: merged_result.xlsx in download_dir)
Returns:
Path to the merged output file
"""
if not downloaded_files:
self._print("No files to post-process")
return None
self._print(f"Post-processing {len(downloaded_files)} downloaded files...")
converter = ExcelConverter(verbose=self.verbose)
all_dfs = []
# Convert each file
for i, file_path in enumerate(downloaded_files):
self._print(f"Converting file {i + 1}/{len(downloaded_files)}: {file_path}")
df = converter.convert(input_file=file_path) # type: ignore
all_dfs.append(df)
# Merge all DataFrames
if all_dfs:
import pandas as pd
merged_df = pd.concat(all_dfs, ignore_index=True)
# Determine output file path
if output_file:
output_path = Path(output_file)
else:
output_path = self.download_dir / "merged_result.xlsx"
# Save merged result
merged_df.to_excel(output_path, index=False)
self._print(f"Merged result saved to: {output_path}")
self._print(f"Total rows: {len(merged_df)}")
return str(output_path)
return None
def extract_and_process(
self, order_ids: List[str], output_file: Optional[str] = None
) -> Optional[str]:
"""
Complete extraction workflow: download + post-process.
Args:
order_ids: List of order IDs to extract
output_file: Path to save merged result (optional)
Returns:
Path to the final merged output file
"""
# Download all batches
downloaded_files = self._download_batches(order_ids)
# Post-process: convert and merge
output_path = self.post_process_downloads(downloaded_files, output_file)
return output_path
def extract_from_file(
self, id_file: str, output_file: Optional[str] = None
) -> Optional[str]:
"""
Extract data from order IDs listed in a file.
Args:
id_file: Path to file containing order IDs (one per line)
output_file: Path to save merged result (optional)
Returns:
Path to the final merged output file
"""
# Read order IDs from file
id_path = Path(id_file)
if not id_path.exists():
raise FileNotFoundError(f"ID file not found: {id_file}")
with open(id_path, "r", encoding="utf-8") as f:
order_ids = [line.strip() for line in f if line.strip()]
self._print(f"Loaded {len(order_ids)} order IDs from {id_file}")
# Extract and process
return self.extract_and_process(order_ids, output_file)
def extract_from_file(
id_file: str,
output_file: Optional[str] = None,
username: Optional[str] = None,
password: Optional[str] = None,
base_url: Optional[str] = None,
headless: bool = True,
ignore_https_errors: bool = True,
verbose: bool = True,
download_dir: Optional[str] = None,
batch_size: int = 10,
) -> Optional[str]:
"""
Convenience function to extract data from order IDs in a file.
Args:
id_file: Path to file containing order IDs (one per line)
output_file: Path to save merged result (optional)
username: Login username (required if not in env)
password: Login password (required if not in env)
base_url: Base URL for ERP system (required if not in env)
headless: Whether to run browser in headless mode
ignore_https_errors: Whether to ignore HTTPS certificate errors
verbose: Whether to print detailed logs
download_dir: Directory to save downloaded files (optional)
batch_size: Number of order IDs per batch
Returns:
Path to the final merged output file
Note:
If username, password, or base_url are not provided,
they will be read from environment variables.
"""
import os
# Get credentials from parameters or environment
if not username:
username = os.getenv("ERP_USERNAME")
if not password:
password = os.getenv("ERP_PASSWORD")
if not base_url:
base_url = os.getenv("ERP_URL")
if not username or not password or not base_url:
raise ValueError(
"username, password, and base_url must be provided either as parameters or environment variables"
)
extractor = DiscreteMaterialPlanExtractor(
username=username,
password=password,
base_url=base_url,
headless=headless,
ignore_https_errors=ignore_https_errors,
verbose=verbose,
download_dir=download_dir,
batch_size=batch_size,
)
try:
return extractor.extract_from_file(id_file, output_file)
finally:
extractor._cleanup_temp_files()