admin/BIPAuto
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(extractor): Add DiscreteMaterialPlanExtractor class for batch data extraction

Browse Source 
- 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
2026-03-27 13:48:07 +08:00
parent bbf335c376
commit da567a2679
2 changed files with 454 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
tests/test_extractor_component.py Normal file
View File

@@ -0,0 +1,120 @@
"""
Test DiscreteMaterialPlanExtractor component structure
"""
import sys
from pathlib import Path
# Add project root to Python path
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
# Load environment variables
from dotenv import load_dotenv
load_dotenv(PROJECT_ROOT / ".env")
# Configure browser path (not used in this test, but consistent with other tests)
import os
os.environ["PLAYWRIGHT_BROWSERS_PATH"] = os.getenv("PLAYWRIGHT_BROWSERS_PATH", "")
print("=" * 60)
print("Testing DiscreteMaterialPlanExtractor Component")
print("=" * 60)
try:
# Test 1: Import the component
print("\n[1/4] Importing DiscreteMaterialPlanExtractor...")
from utils.discrete_material_plan.extractor import (
DiscreteMaterialPlanExtractor,
extract_from_file
)
print("[OK] Import successful")
# Test 2: Verify class exists and has expected methods
print("\n[2/4] Verifying class structure...")
expected_methods = [
'__init__',
'_print',
'_cleanup_temp_files',
'_get_login_url',
'_chunk_order_ids',
'_download_batches',
'post_process_downloads',
'extract_and_process',
'extract_from_file'
]
missing_methods = []
for method in expected_methods:
if not hasattr(DiscreteMaterialPlanExtractor, method):
missing_methods.append(method)
if missing_methods:
print(f"[FAIL] Missing methods: {missing_methods}")
raise AttributeError(f"Missing methods: {missing_methods}")
print(f"[OK] All {len(expected_methods)} expected methods found")
# Test 3: Test instantiation with dummy parameters (no browser)
print("\n[3/4] Testing class instantiation...")
extractor = DiscreteMaterialPlanExtractor(
username="test_user",
password="test_password",
base_url="https://example.com",
headless=True,
ignore_https_errors=True,
verbose=True,
download_dir=None,
batch_size=10
)
# Verify attributes are set correctly
assert extractor.username == "test_user", "Username not set correctly"
assert extractor.password == "test_password", "Password not set correctly"
assert extractor.base_url == "https://example.com", "Base URL not set correctly"
assert extractor.headless == True, "Headless not set correctly"
assert extractor.ignore_https_errors == True, "HTTPS errors setting not set correctly"
assert extractor.verbose == True, "Verbose not set correctly"
assert extractor.batch_size == 10, "Batch size not set correctly"
assert hasattr(extractor, 'download_dir'), "Download dir not set"
print("[OK] Instantiation successful, all attributes verified")
# Test 4: Test helper methods
print("\n[4/4] Testing helper methods...")
# Test _get_login_url
login_url = extractor._get_login_url()
expected_url = "https://example.com/yonbip/resources/uap/rbac/login/main/index.html"
assert login_url == expected_url, f"Login URL incorrect: {login_url}"
print(f"[OK] _get_login_url() returns: {login_url}")
# Test _chunk_order_ids
test_ids = ["ID1", "ID2", "ID3", "ID4", "ID5"]
extractor.batch_size = 2
chunks = extractor._chunk_order_ids(test_ids)
expected_chunks = [["ID1", "ID2"], ["ID3", "ID4"], ["ID5"]]
assert chunks == expected_chunks, f"Chunking incorrect: {chunks}"
print(f"[OK] _chunk_order_ids() works correctly: {chunks}")
# Test _print (should not raise)
extractor._print("Test message")
print("[OK] _print() works correctly")
# Cleanup
extractor._cleanup_temp_files()
print("\n" + "=" * 60)
print("[SUCCESS] All component tests passed!")
print("=" * 60)
except Exception as e:
print(f"\n[ERROR] Component test failed!")
print(f"Error: {e}")
import traceback
traceback.print_exc()
sys.exit(1)
print("\n" + "=" * 60)
print("Test completed successfully")
print("=" * 60)

334
utils/discrete_material_plan/extractor.py Normal file
View 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()