plan.py

import json
from typing import Optional, IO, Tuple, Literal, Union
from .common import request, Result
DateTimeTuple = Tuple[str, str]

def find(mode: Optional[str] = None,
        obs_group: Optional[str] = None,
        proposal_id: Optional[str] = None,
        obs_id: Optional[str] = None,
        instrument: Literal['MSC', 'IFS', 'MCI', 'HSTDM', 'CPIC'] = 'MSC', 
        obs_type: Optional[str] = None,
        object_name: Optional[str] = None,
        obstime: Optional[DateTimeTuple] = None,
        dataset: Optional[str] = None,
        page: int = 1,
        limit: int = 0) -> Result:
    """
    根据给定的参数在DFS中搜索编排数据。
    
    Args:
        mode (Optional[str], optional): 观测模式，如CM1-1、OM1. Defaults to None.
        obs_group (Optional[str], optional): 观测组. Defaults to None.
        proposal_id (Optional[str], optional): 提案号. Defaults to None.
        obs_id (Optional[str], optional): 观测ID，支持模糊搜索. Defaults to None.
        instrument (Optional[str], optional): 模块ID，如'MSC', 'IFS'. Defaults to None.
        obs_type (Optional[str], optional): 观测类型，如主巡天宽场、TOO观测、定标星场等等. Defaults to None.
        object_name (Optional[str], optional): 目标名称. Defaults to None.
        obstime (Optional[DateTimeTuple], optional): 观测时间范围. 如("2021-08-30 00:00:00", "2024-12-30 23:59:59")，Defaults to None. 
        dataset (Optional[str], optional): 数据集名称. Defaults to constants.DEFAULT_DATASET.
        page (int, optional): 页码. Defaults to 1.
        limit (int, optional): 每页数量. Defaults to 0，不限制.
    
    Returns:
        Result: 搜索结果对象.
    
    """

    params = {
        'mode': mode,
        'obs_group': obs_group,
        'proposal_id': proposal_id,
        'obs_id': obs_id,
        'instrument': instrument,
        'obs_type': obs_type,
        'object_name': object_name,
        'obs_time_start': None,
        'obs_time_end': None,
        'create_time_start': None,
        'create_time_end': None,
        'dataset': dataset,
        'page': page,
        'limit': limit,
    }
    
    if obstime is not None:
        params['obs_time_start'], params['obs_time_end'] = obstime
    
    return request.post("/api/plan", params)

def get_by_id(_id: int) -> Result:
    """
    通过计划ID获取编排数据。
    
    Args:
        _id (int): 计划ID。
    
    Returns:
        Result: 查询结果。
    
    """
    return request.get(f"/api/plan/id/{_id}")

def find_by_opid(opid: str) -> Result:
    """
    通过 plan 的 opid 查询编排数据。
    
    Args:
        opid (str): opid。
    
    Returns:
        Result: 查询结果。
    
    """
    return request.get(f"/api/plan/{opid}")


def write_file(local_file: Union[IO, str, list]) -> Result:
    """
    将本地json文件、json数据流或json数据列表写入DFS中。

    Args:
        local_file (Union[IO, str, list]): 文件路径、数据流或JSON数据列表。
        **kwargs: 额外的关键字参数，这些参数将传递给DFS。

    Returns:
        Result: 操作的结果对象，包含操作是否成功以及相关的错误信息，成功返回数据对象。

    """
    batch_size = 512
    
    if isinstance(local_file, str):
        with open(local_file, 'r') as f:
            data = json.load(f)
    elif isinstance(local_file, IO):
        data = json.load(local_file)
    elif isinstance(local_file, list):
        data = local_file
    else:
        raise ValueError("Unsupported type for local_file")

    if not isinstance(data, list):
        raise ValueError("Data must be a list of JSON objects")
    
    for i in range(0, len(data), batch_size):
        batch_data = data[i:i + batch_size]
        response = request.post("/api/plan/file", {"plans": batch_data})
        if not response.success:
            return response
    return Result.ok_data(len(data))

def new(data: dict) -> Result:
    """
    新建编排数据
    
    Args:
        data (dict): 编排数据的字典表示，如：{'id': 394, 'opid': 'xxx', 'instrument': 'MSC', ...}
    
    Returns:
        Result: 成功后，Result.data为写入记录，失败message为失败原因。
    
    """
    return request.post("/api/plan/new", data)

def count_plan_level0(
        instrument: Literal['MSC', 'IFS', 'MCI', 'HSTDM', 'CPIC'], 
        obs_group: str,
        obs_type: str,
        dataset: str,
        obs_id: Optional[str] = None,
        prc_status: Optional[int] = None) -> Result:
    """
    统计编排和0级数据数量
    
    Args:
        instrument (Literal['MSC', 'IFS', 'MCI', 'HSTDM', 'CPIC']): 后端设备，如'MSC', 'IFS'.
        obs_group (str): 观测组.
        obs_type (str): 观测类型，如主巡天宽场、TOO观测、定标星场等等.
        dataset (str): 数据集名称.
        obs_id (Optional[str]): 观测ID. Defaults to None.
        prc_status (Optional[int]): 处理状态. Defaults to None.
    
    Returns:
        Result: 操作的结果对象，包含操作是否成功以及相关的错误信息，成功返回数据对象，如：
        {'code': 200, 'message': '', 'data': {'plan_count': 30, 'level0_count': 30}}
    """
    params = {
        'obs_group': obs_group,
        'instrument': instrument,
        'obs_type': obs_type,
        'dataset': dataset,
        'obs_id': obs_id,
        'prc_status': prc_status
    }    
    return request.post("/api/plan/count", params)