@@ -24,14 +24,30 @@ def get_app():
2424 app = typer .Typer (
2525 name = "heasarc" ,
2626 help = builtins ._ ("Query the HEASARC database." ),
27- invoke_without_command = True ,
28- no_args_is_help = False
27+ invoke_without_command = True , # Allow invoking without a subcommand
28+ # no_args_is_help=True # Show help if no arguments are provided (handled manually below)
2929 )
3030
3131 @app .callback ()
32+ @global_keyboard_interrupt_handler
3233 def heasarc_callback (
3334 ctx : typer .Context ,
34- debug : bool = typer .Option (
35+ # Moved query options to callback
36+ ra : Optional [float ] = typer .Option (None , help = builtins ._ ("Right Ascension in degrees." )),
37+ dec : Optional [float ] = typer .Option (None , help = builtins ._ ("Declination in degrees." )),
38+ radius : Optional [float ] = typer .Option (None , help = builtins ._ ("Search radius in degrees (for cone search)." )),
39+ output_file : Optional [str ] = common_output_options ["output_file" ],
40+ output_format : Optional [str ] = common_output_options ["output_format" ],
41+ max_rows_display : int = typer .Option (
42+ 25 , "--max-rows-display" , help = builtins ._ ("Maximum number of rows to display. Use -1 for all rows." )
43+ ),
44+ show_all_columns : bool = typer .Option (
45+ False , "--show-all-cols" , help = builtins ._ ("Show all columns in the output table." )
46+ ),
47+ max_rows : int = typer .Option (
48+ 100 , "--max-rows" , help = builtins ._ ("Maximum number of rows to retrieve from the HEASARC database. Use -1 for all rows." )
49+ ),
50+ enable_debug : bool = typer .Option (
3551 False ,
3652 "-t" ,
3753 "--debug" ,
@@ -43,64 +59,84 @@ def heasarc_callback(
4359 "-v" ,
4460 "--verbose" ,
4561 help = _ ("Enable verbose output." )
46- )
62+ ),
4763 ):
48- setup_debug_context (ctx , debug , verbose )
49-
50- if ctx .invoked_subcommand is None and \
51- not any (arg in ["-h" , "--help" ] for arg in ctx .args ):
52- help_output_capture = StringIO ()
53- with redirect_stdout (help_output_capture ):
54- try :
55- app (ctx .args + ["--help" ])
56- except SystemExit :
57- pass
58- full_help_text = help_output_capture .getvalue ()
64+ setup_debug_context (ctx , enable_debug , verbose )
65+ debug (f"heasarc_callback: ctx.invoked_subcommand={ ctx .invoked_subcommand } { ctx .args } )
66+ heasarc = Heasarc () # Instantiate Heasarc here for all operations
5967
60- commands_match = re .search (r'╭─ Commands ─.*?(\n(?:│.*?\n)*)╰─.*─╯' , full_help_text , re .DOTALL )
61- if commands_match :
62- commands_section = commands_match .group (0 )
63- filtered_commands_section = "\n " .join ([
64- line for line in commands_section .splitlines () if "Usage:" not in line
65- ])
66- console .print (filtered_commands_section )
67- else :
68- console .print (full_help_text )
69- raise typer .Exit ()
68+ # Check if a subcommand was invoked (e.g., 'list')
69+ if ctx .invoked_subcommand is not None :
70+ return # Let the subcommand handle its own logic
7071
71- @app .command (name = "query" , help = builtins ._ ("Perform a query on HEASARC." ))
72+ @app .command (name = "query" , help = builtins ._ ("Query a specific HEASARC mission. Use 'heasarc list' to see available missions ." ))
7273 @global_keyboard_interrupt_handler
73- def query_heasarc (
74+ def query_heasarc_mission (
7475 ctx : typer .Context ,
75- mission : str = typer .Argument (..., help = builtins ._ ("HEASARC mission name (e.g., 'chandra', 'xmm ')." )),
76+ mission : str = typer .Argument (..., help = builtins ._ ("HEASARC mission name (e.g., 'atnfpsr ')." )),
7677 ra : Optional [float ] = typer .Option (None , help = builtins ._ ("Right Ascension in degrees." )),
7778 dec : Optional [float ] = typer .Option (None , help = builtins ._ ("Declination in degrees." )),
7879 radius : Optional [float ] = typer .Option (None , help = builtins ._ ("Search radius in degrees (for cone search)." )),
7980 output_file : Optional [str ] = common_output_options ["output_file" ],
8081 output_format : Optional [str ] = common_output_options ["output_format" ],
8182 max_rows_display : int = typer .Option (
82- 25 , help = builtins ._ ("Maximum number of rows to display. Use -1 for all rows." )
83+ 25 , "--max-rows-display" , help = builtins ._ ("Maximum number of rows to display. Use -1 for all rows." )
8384 ),
8485 show_all_columns : bool = typer .Option (
8586 False , "--show-all-cols" , help = builtins ._ ("Show all columns in the output table." )
8687 ),
88+ max_rows : int = typer .Option (
89+ 100 , "--max-rows" , help = builtins ._ ("Maximum number of rows to retrieve from the HEASARC database. Use -1 for all rows." )
90+ ),
8791 ):
88- if ctx . obj . get ( "DEBUG" ):
89- debug ( f"query_heasarc - mission: { mission } , ra: { ra } , dec: { dec } , radius: { radius } " )
92+ # Debug context is set up by the heasarc_callback
93+ heasarc = Heasarc () # Instantiate Heasarc here for all operations
9094
9195 console .print (f"[cyan]{ _ ('Querying HEASARC mission: {mission}...' ).format (mission = mission )} )
9296
9397 try :
94- heasarc = Heasarc ()
95- heasarc .mission = mission
98+ # Get the actual catalog name from the mission list
99+ catalogs = heasarc .list_catalogs ()
100+
101+ catalog_name = None
102+ search_mission_upper = mission .strip ().upper ()
103+
104+ # Specific mapping for 'atnfpsr' to 'atnfpulsar'
105+ if search_mission_upper == 'ATNFPSR' :
106+ catalog_name = 'atnfpulsar'
107+ else :
108+ for row in catalogs :
109+ # Use 'name' as the primary key for catalog/table name
110+ if 'name' in row and row ['name' ].strip ().upper () == search_mission_upper :
111+ catalog_name = row ['name' ]
112+ break
113+ # Fallback to 'Table' or 'CATALOG' if 'name' is not present or doesn't match
114+ elif 'Table' in row and row ['Table' ].strip ().upper () == search_mission_upper :
115+ catalog_name = row ['Table' ]
116+ break
117+ elif 'CATALOG' in row and row ['CATALOG' ].strip ().upper () == search_mission_upper :
118+ catalog_name = row ['CATALOG' ]
119+ break
120+
121+ if catalog_name is None :
122+ console .print (_ (f"[red]Error: Mission '{ mission } ))
123+ raise typer .Exit (code = 1 )
124+
125+ # Set maxrec for query
126+ maxrec_value = None if max_rows == - 1 else max_rows
96127
97128 if ra is not None and dec is not None :
98129 if radius is None :
99130 console .print (_ ("[red]Error: --radius is required when --ra and --dec are provided for a cone search.[/red]" ))
100131 raise typer .Exit (code = 1 )
101- results = heasarc .query_region (ra = ra , dec = dec , radius = radius )
132+ adql_query = heasarc .query_region (ra = ra , dec = dec , radius = radius , catalog = catalog_name , get_query_payload = True , maxrec = maxrec_value )
133+ console .print (f"[debug] ADQL Query (cone): { adql_query } )
134+ results = heasarc .query_region (ra = ra , dec = dec , radius = radius , catalog = catalog_name , maxrec = maxrec_value )
102135 else :
103- results = heasarc .query_object (mission ) # Query all data for the mission
136+ # Use query_tap for all-sky queries with the correct catalog name
137+ adql_query_string = f"SELECT * FROM { catalog_name }
138+ console .print (f"[debug] ADQL Query (all-sky): { adql_query_string } )
139+ results = heasarc .query_tap (query = adql_query_string , maxrec = maxrec_value )
104140
105141 if results and len (results ) > 0 :
106142 console .print (_ ("[green]Found {count} result(s) from HEASARC.[/green]" ).format (count = len (results )))
@@ -114,20 +150,52 @@ def query_heasarc(
114150 handle_astroquery_exception (ctx , e , _ ("HEASARC query" ))
115151 raise typer .Exit (code = 1 )
116152
117- @app .command (name = "list-missions " , help = builtins ._ ("List available HEASARC missions." ))
153+ @app .command (name = "list" , help = builtins ._ ("List available HEASARC missions." ))
118154 @global_keyboard_interrupt_handler
119- def list_missions (ctx : typer .Context ):
155+ def list_heasarc_missions (
156+ ctx : typer .Context ,
157+ max_rows_display : int = typer .Option (
158+ 25 , "--max-rows-display" , "--max-rows" , help = builtins ._ ("Maximum number of rows to display. Use -1 for all rows." )
159+ ),
160+ show_all_columns : bool = typer .Option (
161+ False , "--show-all-cols" , help = builtins ._ ("Show all columns in the output table." )
162+ ),
163+ prefix : Optional [str ] = typer .Option (
164+ None , "--prefix" , "-p" , help = builtins ._ ("Filter missions by a starting prefix (e.g., 'atnfpsr')." )
165+ ),
166+ ):
167+ debug (f"list_heasarc_missions: max_rows_display={ max_rows_display } { show_all_columns } { prefix } )
168+ heasarc = Heasarc ()
120169 console .print (f"[cyan]{ _ ('Listing HEASARC missions...' )} )
121170 try :
122- missions = Heasarc .get_missions ()
123- if missions :
124- console .print (_ ("[green]Available HEASARC Missions:[/green]" ))
125- for m in missions :
126- console .print (f"- { m } )
171+ missions_table = heasarc .list_catalogs ()
172+
173+ if prefix :
174+ original_count = len (missions_table )
175+ # Determine the correct column name for filtering
176+ mission_name_column = None
177+ if 'name' in missions_table .colnames :
178+ mission_name_column = 'name'
179+ elif 'Table' in missions_table .colnames :
180+ mission_name_column = 'Table'
181+ elif 'CATALOG' in missions_table .colnames :
182+ mission_name_column = 'CATALOG'
183+
184+ if mission_name_column :
185+ # Apply the filter using boolean indexing
186+ mask = [str (name ).lower ().startswith (prefix .lower ()) for name in missions_table [mission_name_column ]]
187+ missions_table = missions_table [mask ]
188+ else :
189+ # If no suitable column is found, log a warning or handle appropriately
190+ debug ("Could not find a suitable mission name column ('name', 'Table', or 'CATALOG') for filtering." )
191+
192+ console .print (_ ("[cyan]Filtered from {original_count} to {filtered_count} missions with prefix \" {prefix}\" .[/cyan]" ).format (original_count = original_count , filtered_count = len (missions_table ), prefix = prefix ))
193+
194+ if missions_table :
195+ display_table (ctx , missions_table , title = _ ("Available HEASARC Missions" ), max_rows = max_rows_display , show_all_columns = show_all_columns )
127196 else :
128- console .print (_ ("[yellow]No HEASARC missions found.[/yellow]" ))
197+ console .print (_ ("[yellow]No HEASARC missions found matching your criteria .[/yellow]" ))
129198 except Exception as e :
130199 handle_astroquery_exception (ctx , e , _ ("HEASARC list missions" ))
131200 raise typer .Exit (code = 1 )
132-
133201 return app
0 commit comments