"""
A grammar taking rows from a FITS table.
"""

#c Copyright 2008-2022, the GAVO project <gavo@ari.uni-heidelberg.de>
#c
#c This program is free software, covered by the GNU GPL.  See the
#c COPYING file in the source distribution.


from gavo import base
from gavo.grammars import common

from gavo.utils import pyfits


class FITSTableIterator(common.RowIterator):
	"""The row iterator for FITSTableGrammars.
	"""
	def _iterRows(self):
		hdus = pyfits.open(self.sourceToken)
		fitsTable = hdus[self.grammar.hdu].data
		names = [n for n in fitsTable.dtype.names]

		if self.grammar.lowerKeys:
			names = [n.lower() for n in names]

		for row in fitsTable:
			if self.grammar.nanIsNULL:
				row = [v if v==v else None for v in row]

			res = dict(list(zip(names, row)))
			yield res


class FITSTableGrammar(common.Grammar):
	"""A grammar parsing from FITS tables.

	fitsTableGrammar result in typed records, i.e., values normally come
	in the types they are supposed to have.  Of course, that won't work
	for datetimes, STC-S regions, and the like.

	The keys of the result dictionaries are simpily the names given in
	the FITS.
	"""
	name_ = "fitsTableGrammar"

	_hduIndex = base.IntAttribute("hdu", default=1,
		description="Take the data from this extension (primary=0)."
		" Tabular data typically resides in the first extension.")
	_lowerKeys = base.BooleanAttribute("lowerKeys", default=False,
		description="Convert column names to lower case when building"
		" the rawdict.")
	_mapNaN = base.BooleanAttribute("nanIsNULL", default=False,
		description="Map NaN floats to NULL when building the rawdict.")

	rowIterator = FITSTableIterator
