catalog.py

from .common import request, Result
import pickle

def search(ra: float,
        dec: float,
        radius: float,
        catalog_name: str,
        columns: list[str] = ['*'],
        min_mag: float = -1,
        max_mag: float = -1,
        obstime: int = 0,
        limit: int = 0) -> Result:
    """
    根据输入的赤经、赤纬、半径、星表名称、列名、最小星等、最大星等、观测时间和限制数量，
    查询指定范围内的天体数据，并返回查询结果。

    Args:
        ra (float): 天体赤经（以度为单位）。
        dec (float): 天体赤纬（以度为单位）。
        radius (float): 查询半径（以度为单位）。
        catalog_name (str): 星表名称。
        columns (list[str], optional): 查询的列名，默认为 ['*']。表示查询所有列。
        min_mag (float, optional): 最小星等，默认为 -1。表示无限制。
        max_mag (float, optional): 最大星等，默认为 -1。表示无限制。
        obstime (int, optional): 观测时间，默认为 0，暂无用。
        limit (int, optional): 查询结果的数量限制，默认为 0。表示无限制。

    Returns:
        Result: 查询结果。

    Raises:
        Exception: 如果请求失败，将抛出异常。

    """

    if not columns:
        columns = ['*']

    try:
        import pandas as pd
        if not catalog_name:
            return Result.error("catalog_name is required")
        datas = request.post('/api/catalog', {
                'ra': ra,
                'dec': dec,
                'radius': radius,
                'catalog_name': catalog_name,
                'columns': columns,
                'min_mag': min_mag,
                'max_mag': max_mag,
                'obstime': obstime,
                'limit': limit
            })
        if datas and isinstance(datas, Result):
            return datas
        records = pickle.loads(datas._content)
        df, total_count = records['records'], records['totalCount']
        return Result.ok_data(data = df).append("totalCount", total_count)
    except Exception as e:
        return Result.error(str(e))


def all_catalog_names() -> Result:
    """
    获取星表名称列表。

    Returns:
        Result: 星表列表。

    Raises:
        Exception: 如果请求失败，将抛出异常。

    """
    try:
        datas = request.get('/api/catalog/names')
        if datas and isinstance(datas, Result):
            return datas
        records = pickle.loads(datas._content)
        return Result.ok_data(data = records)
    except ImportError:
        return Result.error("pandas is not installed")        
    except Exception as e:
        return Result.error(str(e))

def describe_catalog(catalog_name: str) -> Result:
    """
    获取星表各列信息。

    Returns:
        Result: 列描述的列表。

    Raises:
        Exception: 如果请求失败，将抛出异常。

    """
    try:
        import pandas as pd
        if not catalog_name:
            return Result.error("catalog_name is required")
        datas = request.get(f'/api/catalog/describe/{catalog_name}')
        if datas and isinstance(datas, Result):
            return datas
        records = pickle.loads(datas._content)
        return Result.ok_data(data = records)
    except ImportError:
        return Result.error("pandas is not installed")    
    except Exception as e:
        return Result.error(str(e))