regularize formatting of list_majorbodies and list_sites

lhorizon/handlers.py

@@ -172,6 +172,15 @@ def query_all_lhorizons(
             time.sleep(delay_between)
 
 
+def _format_site_id(obj):
+    if isinstance(obj, float) and np.isnan(obj):
+        return None
+    try:
+        return int(obj)
+    except (ValueError, TypeError):
+        return obj
+
+
 def list_sites(center_body: int = 399) -> pd.DataFrame:
     """
     query _Horizons_ for all named sites recognized on the specified body and
@@ -197,7 +206,10 @@ def list_sites(center_body: int = 399) -> pd.DataFrame:
         if "." in str(split[0]):
             split = [np.nan] + re.split(" +", line.strip(), maxsplit=4)
         data_values.append(split)
-    return pd.DataFrame(data_values, columns=columns)
+    df = pd.DataFrame(data_values, columns=columns)
+    df = df.rename(columns={'code_subc': 'id', 'Site name': 'name'})
+    df['id'] = df['id'].map(_format_site_id)
+    return df
 
 
 def list_majorbodies() -> pd.DataFrame:
@@ -217,7 +229,9 @@ def list_majorbodies() -> pd.DataFrame:
         strip = re.sub(" +", " ", line.strip())
         split = re.split(" +", strip, maxsplit=1)
         data_values.append(split)
-    return pd.DataFrame(data_values, columns=["id", "name"]).dropna(axis=0)
+    df = pd.DataFrame(data_values, columns=["id", "name"]).dropna(axis=0)
+    df['id'] = df['id'].astype(int)
+    return df
 
 
 def get_observer_quantity_codes() -> str:

lhorizon/tests/test_handlers.py

@@ -70,9 +70,12 @@ def test_list_sites():
     """
     sites = list_sites()
     assert (
-        sites.loc[sites["code_subc"] == "-1"]["Site name"].iloc[0] == "Arecibo (300-m, 1963 to 2020)"
+        sites.loc[sites["id"] == -1]["name"].iloc[0]
+        == "Arecibo (300-m, 1963 to 2020)"
+    )
+    assert ",".join(sites.columns) == (
+        "id,east long,cyl_rho,z,epoch,name"
     )
-    assert ",".join(sites.columns) == "code_subc,east long,cyl_rho,z,epoch,Site name"
 
 
 def test_list_majorbodies():
@@ -82,7 +85,7 @@ def test_list_majorbodies():
     """
     bodies = list_majorbodies()
     assert (
-        bodies.loc[bodies["id"] == "1"]["name"].iloc[0] == "Mercury Barycenter"
+        bodies.loc[bodies["id"] == 1]["name"].iloc[0] == "Mercury Barycenter"
     )
     assert all(
         [