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
This commit is contained in:
Misaka_Company
334
utils/discrete_material_plan/extractor.py
Normal file
334
utils/discrete_material_plan/extractor.py
Normal file
@@ -0,0 +1,334 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user