1# Licensed under a 3-clause BSD style license - see LICENSE.rst
2"""
3This module's main purpose is to act as a script to create new versions
4of ufunc.c when ERFA is updated (or this generator is enhanced).
5
6`Jinja2 <http://jinja.pocoo.org/>`_ must be installed for this
7module/script to function.
8
9Note that this does *not* currently automate the process of creating structs
10or dtypes for those structs.  They should be added manually in the template file.
11"""
12
13import re
14import os.path
15from collections import OrderedDict
16
17DEFAULT_ERFA_LOC = os.path.join(os.path.split(__file__)[0], 'liberfa/erfa/src')
18DEFAULT_TEMPLATE_LOC = os.path.join(os.path.split(__file__)[0], 'erfa')
19
20NDIMS_REX = re.compile(re.escape("numpy.dtype([('fi0', '.*', <(.*)>)])")
21                       .replace(r'\.\*', '.*')
22                       .replace(r'\<', '(')
23                       .replace(r'\>', ')'))
24
25
26class FunctionDoc:
27
28    def __init__(self, doc):
29        self.doc = doc.replace("**", "      ").replace("/*\n", "").replace("*/", "")
30        self.doc = self.doc.replace("/*+\n", "")        # accommodate eraLdn
31        self.doc = self.doc.replace("*  ", "    " * 2)  # accommodate eraAticqn
32        self.doc = self.doc.replace("*\n", "\n")        # accommodate eraAticqn
33        self.__input = None
34        self.__output = None
35        self.__ret_info = None
36
37    def _get_arg_doc_list(self, doc_lines):
38        """Parse input/output doc section lines, getting arguments from them.
39
40        Ensure all elements of eraASTROM and eraLDBODY are left out, as those
41        are not input or output arguments themselves.  Also remove the nb
42        argument in from of eraLDBODY, as we infer nb from the python array.
43        """
44        doc_list = []
45        skip = []
46        for d in doc_lines:
47            arg_doc = ArgumentDoc(d)
48            if arg_doc.name is not None:
49                if skip:
50                    if skip[0] == arg_doc.name:
51                        skip.pop(0)
52                        continue
53                    else:
54                        raise RuntimeError("We whould be skipping {} "
55                                           "but {} encountered."
56                                           .format(skip[0], arg_doc.name))
57
58                if arg_doc.type.startswith('eraLDBODY'):
59                    # Special-case LDBODY: for those, the previous argument
60                    # is always the number of bodies, but we don't need it
61                    # as an input argument for the ufunc since we're going
62                    # to determine this from the array itself. Also skip
63                    # the description of its contents; those are not arguments.
64                    doc_list.pop()
65                    skip = ['bm', 'dl', 'pv']
66                elif arg_doc.type.startswith('eraASTROM'):
67                    # Special-case ASTROM: need to skip the description
68                    # of its contents; those are not arguments.
69                    skip = ['pmt', 'eb', 'eh', 'em', 'v', 'bm1',
70                            'bpn', 'along', 'xpl', 'ypl', 'sphi',
71                            'cphi', 'diurab', 'eral', 'refa', 'refb']
72
73                doc_list.append(arg_doc)
74
75        return doc_list
76
77    @property
78    def input(self):
79        if self.__input is None:
80            self.__input = []
81            for regex in ("Given([^\n]*):.*?\n(.+?)  \n",
82                          "Given and returned([^\n]*):\n(.+?)  \n"):
83                result = re.search(regex, self.doc, re.DOTALL)
84                if result is not None:
85                    doc_lines = result.group(2).split("\n")
86                    self.__input += self._get_arg_doc_list(doc_lines)
87
88        return self.__input
89
90    @property
91    def output(self):
92        if self.__output is None:
93            self.__output = []
94            for regex in ("Given and returned([^\n]*):\n(.+?)  \n",
95                          "Returned([^\n]*):.*?\n(.+?)  \n"):
96                result = re.search(regex, self.doc, re.DOTALL)
97                if result is not None:
98                    doc_lines = result.group(2).split("\n")
99                    self.__output += self._get_arg_doc_list(doc_lines)
100
101        return self.__output
102
103    @property
104    def ret_info(self):
105        if self.__ret_info is None:
106            ret_info = []
107            result = re.search("Returned \\(function value\\)([^\n]*):\n(.+?)  \n",
108                               self.doc, re.DOTALL)
109            if result is not None:
110                ret_info.append(ReturnDoc(result.group(2)))
111
112            if len(ret_info) == 0:
113                self.__ret_info = ''
114            elif len(ret_info) == 1:
115                self.__ret_info = ret_info[0]
116            else:
117                raise ValueError("Multiple C return sections found in this doc:\n"
118                                 + self.doc)
119
120        return self.__ret_info
121
122    @property
123    def title(self):
124        # Used for the docstring title.
125        lines = [line.strip() for line in self.doc.split('\n')[4:10]]
126        # Always include the first line, then stop at either an empty
127        # line or at the end of a sentence.
128        description = lines[:1]
129        for line in lines[1:]:
130            if line == '':
131                break
132            if '. ' in line:
133                line = line[:line.index('. ')+1]
134            description.append(line)
135            if line.endswith('.'):
136                break
137
138        return '\n    '.join(description)
139
140    def __repr__(self):
141        return '\n'.join([(ln.rstrip() if ln.strip() else '')
142                          for ln in self.doc.split('\n')])
143
144
145class ArgumentDoc:
146
147    def __init__(self, doc):
148        match = re.search("^ +([^ ]+)[ ]+([^ ]+)[ ]+(.+)", doc)
149        if match is not None:
150            self.name = match.group(1)
151            if self.name.startswith('*'):  # Easier than getting the regex to behave...
152                self.name = self.name.replace('*', '')
153            self.type = match.group(2)
154            self.doc = match.group(3)
155        else:
156            self.name = None
157            self.type = None
158            self.doc = None
159
160    def __repr__(self):
161        return f"    {self.name:15} {self.type:15} {self.doc}"
162
163
164class Variable:
165    """Properties shared by Argument and Return."""
166    @property
167    def npy_type(self):
168        """Predefined type used by numpy ufuncs to indicate a given ctype.
169
170        Eg., NPY_DOUBLE for double.
171        """
172        return "NPY_" + self.ctype.upper()
173
174    @property
175    def dtype(self):
176        """Name of dtype corresponding to the ctype.
177
178        Specifically,
179        double : dt_double
180        int : dt_int
181        double[3]: dt_vector
182        double[2][3] : dt_pv
183        double[2] : dt_pvdpv
184        double[3][3] : dt_matrix
185        int[4] : dt_ymdf | dt_hmsf | dt_dmsf, depding on name
186        eraASTROM: dt_eraASTROM
187        eraLDBODY: dt_eraLDBODY
188        char : dt_sign
189        char[] : dt_type
190
191        The corresponding dtypes are defined in ufunc.c, where they are
192        used for the loop definitions.  In core.py, they are also used
193        to view-cast regular arrays to these structured dtypes.
194        """
195        if self.ctype == 'const char':
196            return 'dt_type'
197        elif self.ctype == 'char':
198            return 'dt_sign'
199        elif self.ctype == 'int' and self.shape == (4,):
200            return 'dt_' + self.name[1:]
201        elif self.ctype == 'double' and self.shape == (3,):
202            return 'dt_double'
203        elif self.ctype == 'double' and self.shape == (2, 3):
204            return 'dt_pv'
205        elif self.ctype == 'double' and self.shape == (2,):
206            return 'dt_pvdpv'
207        elif self.ctype == 'double' and self.shape == (3, 3):
208            return 'dt_double'
209        elif not self.shape:
210            return 'dt_' + self.ctype
211        else:
212            raise ValueError("ctype {} with shape {} not recognized."
213                             .format(self.ctype, self.shape))
214
215    @property
216    def view_dtype(self):
217        """Name of dtype corresponding to the ctype for viewing back as array.
218
219        E.g., dt_double for double, dt_double33 for double[3][3].
220
221        The types are defined in core.py, where they are used for view-casts
222        of structured results as regular arrays.
223        """
224        if self.ctype == 'const char':
225            return 'dt_bytes12'
226        elif self.ctype == 'char':
227            return 'dt_bytes1'
228        else:
229            raise ValueError('Only char ctype should need view back!')
230
231    @property
232    def ndim(self):
233        return len(self.shape)
234
235    @property
236    def size(self):
237        size = 1
238        for s in self.shape:
239            size *= s
240        return size
241
242    @property
243    def cshape(self):
244        return ''.join([f'[{s}]' for s in self.shape])
245
246    @property
247    def signature_shape(self):
248        if self.ctype == 'eraLDBODY':
249            return '(n)'
250        elif self.ctype == 'double' and self.shape == (3,):
251            return '(3)'
252        elif self.ctype == 'double' and self.shape == (3, 3):
253            return '(3, 3)'
254        else:
255            return '()'
256
257
258class Argument(Variable):
259
260    def __init__(self, definition, doc):
261        self.definition = definition
262        self.doc = doc
263        self.__inout_state = None
264        self.ctype, ptr_name_arr = definition.strip().rsplit(" ", 1)
265        if "*" == ptr_name_arr[0]:
266            self.is_ptr = True
267            name_arr = ptr_name_arr[1:]
268        else:
269            self.is_ptr = False
270            name_arr = ptr_name_arr
271        if "[]" in ptr_name_arr:
272            self.is_ptr = True
273            name_arr = name_arr[:-2]
274        if "[" in name_arr:
275            self.name, arr = name_arr.split("[", 1)
276            self.shape = tuple([int(size) for size in arr[:-1].split("][")])
277        else:
278            self.name = name_arr
279            self.shape = ()
280
281    @property
282    def inout_state(self):
283        if self.__inout_state is None:
284            self.__inout_state = ''
285            for i in self.doc.input:
286                if self.name in i.name.split(','):
287                    self.__inout_state = 'in'
288            for o in self.doc.output:
289                if self.name in o.name.split(','):
290                    if self.__inout_state == 'in':
291                        self.__inout_state = 'inout'
292                    else:
293                        self.__inout_state = 'out'
294        return self.__inout_state
295
296    @property
297    def name_for_call(self):
298        """How the argument should be used in the call to the ERFA function.
299
300        This takes care of ensuring that inputs are passed by value,
301        as well as adding back the number of bodies for any LDBODY argument.
302        The latter presumes that in the ufunc inner loops, that number is
303        called 'nb'.
304        """
305        if self.ctype == 'eraLDBODY':
306            assert self.name == 'b'
307            return 'nb, _' + self.name
308        elif self.is_ptr:
309            return '_'+self.name
310        else:
311            return '*_'+self.name
312
313    def __repr__(self):
314        return (f"Argument('{self.definition}', name='{self.name}', "
315                f"ctype='{self.ctype}', inout_state='{self.inout_state}')")
316
317
318class ReturnDoc:
319
320    def __init__(self, doc):
321        self.doc = doc
322
323        self.infoline = doc.split('\n')[0].strip()
324        self.type = self.infoline.split()[0]
325        self.descr = self.infoline.split()[1]
326
327        if self.descr.startswith('status'):
328            self.statuscodes = statuscodes = {}
329
330            code = None
331            for line in doc[doc.index(':')+1:].split('\n'):
332                ls = line.strip()
333                if ls != '':
334                    if ' = ' in ls:
335                        code, msg = ls.split(' = ')
336                        if code != 'else':
337                            code = int(code)
338                        statuscodes[code] = msg
339                    elif code is not None:
340                        statuscodes[code] += ls
341        else:
342            self.statuscodes = None
343
344    def __repr__(self):
345        return f"Return value, type={self.type:15}, {self.descr}, {self.doc}"
346
347
348class Return(Variable):
349
350    def __init__(self, ctype, doc):
351        self.name = 'c_retval'
352        self.inout_state = 'stat' if ctype == 'int' else 'ret'
353        self.ctype = ctype
354        self.shape = ()
355        self.doc = doc
356
357    def __repr__(self):
358        return f"Return(name='{self.name}', ctype='{self.ctype}', inout_state='{self.inout_state}')"
359
360    @property
361    def doc_info(self):
362        return self.doc.ret_info
363
364
365class Function:
366    """
367    A class representing a C function.
368
369    Parameters
370    ----------
371    name : str
372        The name of the function
373    source_path : str
374        Either a directory, which means look for the function in a
375        stand-alone file (like for the standard ERFA distribution), or a
376        file, which means look for the function in that file.
377    match_line : str, optional
378        If given, searching of the source file will skip until it finds
379        a line matching this string, and start from there.
380    """
381
382    def __init__(self, name, source_path, match_line=None):
383        self.name = name
384        self.pyname = name.split('era')[-1].lower()
385        self.filename = self.pyname+".c"
386        if os.path.isdir(source_path):
387            self.filepath = os.path.join(os.path.normpath(source_path), self.filename)
388        else:
389            self.filepath = source_path
390
391        with open(self.filepath) as f:
392            if match_line:
393                line = f.readline()
394                while line != '':
395                    if line.startswith(match_line):
396                        filecontents = '\n' + line + f.read()
397                        break
398                    line = f.readline()
399                else:
400                    msg = ('Could not find the match_line "{0}" in '
401                           'the source file "{1}"')
402                    raise ValueError(msg.format(match_line, self.filepath))
403            else:
404                filecontents = f.read()
405
406        pattern = fr"\n([^\n]+{name} ?\([^)]+\)).+?(/\*.+?\*/)"
407        p = re.compile(pattern, flags=re.DOTALL | re.MULTILINE)
408
409        search = p.search(filecontents)
410        self.cfunc = " ".join(search.group(1).split())
411        self.doc = FunctionDoc(search.group(2))
412
413        self.args = []
414        for arg in re.search(r"\(([^)]+)\)", self.cfunc).group(1).split(', '):
415            self.args.append(Argument(arg, self.doc))
416        self.ret = re.search(f"^(.*){name}", self.cfunc).group(1).strip()
417        if self.ret != 'void':
418            self.args.append(Return(self.ret, self.doc))
419
420    def args_by_inout(self, inout_filter, prop=None, join=None):
421        """
422        Gives all of the arguments and/or returned values, depending on whether
423        they are inputs, outputs, etc.
424
425        The value for `inout_filter` should be a string containing anything
426        that arguments' `inout_state` attribute produces.  Currently, that can be:
427
428          * "in" : input
429          * "out" : output
430          * "inout" : something that's could be input or output (e.g. a struct)
431          * "ret" : the return value of the C function
432          * "stat" : the return value of the C function if it is a status code
433
434        It can also be a "|"-separated string giving inout states to OR
435        together.
436        """
437        result = []
438        for arg in self.args:
439            if arg.inout_state in inout_filter.split('|'):
440                if prop is None:
441                    result.append(arg)
442                else:
443                    result.append(getattr(arg, prop))
444        if join is not None:
445            return join.join(result)
446        else:
447            return result
448
449    @property
450    def user_dtype(self):
451        """The non-standard dtype, if any, needed by this function's ufunc.
452
453        This would be any structured array for any input or output, but
454        we give preference to LDBODY, since that also decides that the ufunc
455        should be a generalized ufunc.
456        """
457        user_dtype = None
458        for arg in self.args_by_inout('in|inout|out'):
459            if arg.ctype == 'eraLDBODY':
460                return arg.dtype
461            elif user_dtype is None and arg.dtype not in ('dt_double',
462                                                          'dt_int'):
463                user_dtype = arg.dtype
464
465        return user_dtype
466
467    @property
468    def signature(self):
469        """Possible signature, if this function should be a gufunc."""
470        if all(arg.signature_shape == '()'
471               for arg in self.args_by_inout('in|inout|out')):
472            return None
473
474        return '->'.join(
475            [','.join([arg.signature_shape for arg in args])
476             for args in (self.args_by_inout('in|inout'),
477                          self.args_by_inout('inout|out|ret|stat'))])
478
479    @property
480    def python_call(self):
481        out = ', '.join([arg.name for arg in self.args_by_inout('inout|out|stat|ret')])
482        args = ', '.join([arg.name for arg in self.args_by_inout('in|inout')])
483        result = '{out} = {func}({args})'.format(out=out,
484                                                 func='ufunc.' + self.pyname,
485                                                 args=args)
486        if len(result) < 75:
487            return result
488
489        if result.index('(') < 75:
490            return result.replace('(', '(\n        ')
491
492        split_point = result[:75].rfind(',') + 1
493        return ('(' + result[:split_point] + '\n    '
494                + result[split_point:].replace(' =', ') ='))
495
496    def __repr__(self):
497        return (f"Function(name='{self.name}', pyname='{self.pyname}', "
498                f"filename='{self.filename}', filepath='{self.filepath}')")
499
500
501class Constant:
502
503    def __init__(self, name, value, doc):
504        self.name = name.replace("ERFA_", "")
505        self.value = value.replace("ERFA_", "")
506        self.doc = doc
507
508
509class ExtraFunction(Function):
510    """
511    An "extra" function - e.g. one not following the SOFA/ERFA standard format.
512
513    Parameters
514    ----------
515    cname : str
516        The name of the function in C
517    prototype : str
518        The prototype for the function (usually derived from the header)
519    pathfordoc : str
520        The path to a file that contains the prototype, with the documentation
521        as a multiline string *before* it.
522    """
523
524    def __init__(self, cname, prototype, pathfordoc):
525        self.name = cname
526        self.pyname = cname.split('era')[-1].lower()
527        self.filepath, self.filename = os.path.split(pathfordoc)
528
529        self.prototype = prototype.strip()
530        if prototype.endswith('{') or prototype.endswith(';'):
531            self.prototype = prototype[:-1].strip()
532
533        incomment = False
534        lastcomment = None
535        with open(pathfordoc, 'r') as f:
536            for ln in f:
537                if incomment:
538                    if ln.lstrip().startswith('*/'):
539                        incomment = False
540                        lastcomment = ''.join(lastcomment)
541                    else:
542                        if ln.startswith('**'):
543                            ln = ln[2:]
544                        lastcomment.append(ln)
545                else:
546                    if ln.lstrip().startswith('/*'):
547                        incomment = True
548                        lastcomment = []
549                    if ln.startswith(self.prototype):
550                        self.doc = lastcomment
551                        break
552            else:
553                raise ValueError('Did not find prototype {} in file '
554                                 '{}'.format(self.prototype, pathfordoc))
555
556        self.args = []
557        argset = re.search(fr"{self.name}\(([^)]+)?\)",
558                           self.prototype).group(1)
559        if argset is not None:
560            for arg in argset.split(', '):
561                self.args.append(Argument(arg, self.doc))
562        self.ret = re.match(f"^(.*){self.name}",
563                            self.prototype).group(1).strip()
564        if self.ret != 'void':
565            self.args.append(Return(self.ret, self.doc))
566
567    def __repr__(self):
568        r = super().__repr__()
569        if r.startswith('Function'):
570            r = 'Extra' + r
571        return r
572
573
574class TestFunction:
575    """Function holding information about a test in t_erfa_c.c"""
576    def __init__(self, name, t_erfa_c, nin, ninout, nout):
577        self.name = name
578        # Get lines that test the given erfa function: capture everything
579        # between a line starting with '{' after the test function definition
580        # and the first line starting with '}' or ' }'.
581        pattern = fr"\nstatic void t_{name}\(" + r".+?(^\{.+?^\s?\})"
582        search = re.search(pattern, t_erfa_c, flags=re.DOTALL | re.MULTILINE)
583        self.lines = search.group(1).split('\n')
584        # Number of input, inplace, and output arguments.
585        self.nin = nin
586        self.ninout = ninout
587        self.nout = nout
588        # Dict of dtypes for variables, filled by define_arrays().
589        self.var_dtypes = {}
590
591    @classmethod
592    def from_function(cls, func, t_erfa_c):
593        """Initialize from a function definition."""
594        return cls(name=func.pyname, t_erfa_c=t_erfa_c,
595                   nin=len(func.args_by_inout('in')),
596                   ninout=len(func.args_by_inout('inout')),
597                   nout=len(func.args_by_inout('out')))
598
599    def xfail(self):
600        """Whether the python test produced for this function will fail.
601
602        Right now this will be true for functions without inputs such
603        as eraIr.
604        """
605        if self.nin + self.ninout == 0:
606            if self.name == 'zpv':
607                # Works on newer numpy
608                return "np.__version__ < '1.21', reason='needs numpy >= 1.21'"
609            else:
610                return "reason='do not yet support no-input ufuncs'"
611        else:
612            return None
613
614    def pre_process_lines(self):
615        """Basic pre-processing.
616
617        Combine multi-part lines, strip braces, semi-colons, empty lines.
618        """
619        lines = []
620        line = ''
621        for part in self.lines:
622            part = part.strip()
623            if part in ('', '{', '}'):
624                continue
625            line += part + ' '
626            if part.endswith(';'):
627                lines.append(line.strip()[:-1])
628                line = ''
629        return lines
630
631    def define_arrays(self, line):
632        """Check variable definition line for items also needed in python.
633
634        E.g., creating an empty astrom structured array.
635        """
636        defines = []
637        # Split line in type and variables.
638        # E.g., "double x, y, z" will give ctype='double; variables='x, y, z'
639        ctype, _, variables = line.partition(' ')
640        for var in variables.split(','):
641            var = var.strip()
642            # Is variable an array?
643            name, _, rest = var.partition('[')
644            # If not, or one of iymdf or ihmsf, ignore (latter are outputs only).
645            if not rest or rest[:2] == '4]':
646                continue
647            if ctype == 'eraLDBODY':
648                # Special case, since this should be recarray for access similar
649                # to C struct.
650                v_dtype = 'dt_eraLDBODY'
651                v_shape = rest[:rest.index(']')]
652                extra = ".view(np.recarray)"
653            else:
654                # Temporarily create an Argument, so we can use its attributes.
655                # This translates, e.g., double pv[2][3] to dtype dt_pv.
656                v = Argument(ctype + ' ' + var.strip(), '')
657                v_dtype = v.dtype
658                v_shape = v.shape if v.signature_shape != '()' else '()'
659                extra = ""
660            self.var_dtypes[name] = v_dtype
661            if v_dtype == 'dt_double':
662                v_dtype = 'float'
663            else:
664                v_dtype = 'erfa_ufunc.' + v_dtype
665            defines.append(f"{name} = np.empty({v_shape}, {v_dtype}){extra}")
666
667        return defines
668
669    def to_python(self):
670        """Lines defining the body of a python version of the test function."""
671        # TODO: this is quite hacky right now!  Would be good to let function
672        # calls be understood by the Function class.
673
674        # Name of the erfa C function, so that we can recognize it.
675        era_name = 'era' + self.name.capitalize()
676        # Collect actual code lines, without ";", braces, etc.
677        lines = self.pre_process_lines()
678        out = []
679        for line in lines:
680            # In ldn ufunc, the number of bodies is inferred from the array size,
681            # so no need to keep the definition.
682            if line == 'n = 3' and self.name == 'ldn':
683                continue
684
685            # Are we dealing with a variable definition that also sets it?
686            # (hack: only happens for double).
687            if line.startswith('double') and '=' in line:
688                # Complete hack for single occurrence.
689                if line.startswith('double xyz[] = {'):
690                    out.append(f"xyz = np.array([{line[16:-1]}])")
691                else:
692                    # Put each definition on a separate line.
693                    out.extend([part.strip() for part in line[7:].split(',')])
694                continue
695
696            # Variable definitions: add empty array definition as needed.
697            if line.startswith(('double', 'int', 'char', 'eraASTROM', 'eraLDBODY')):
698                out.extend(self.define_arrays(line))
699                continue
700
701            # Actual function. Start with basic replacements.
702            line = (line
703                    .replace('ERFA_', 'erfa.')
704                    .replace('(void)', '')
705                    .replace('(int)', '')
706                    .replace("pv[0]", "pv['p']")
707                    .replace("pv[1]", "pv['v']")
708                    .replace("s, '-'", "s[0], b'-'")  # Rather hacky...
709                    .replace("s, '+'", "s[0], b'+'")  # Rather hacky...
710                    .strip())
711
712            # Call of test function vvi or vvd.
713            if line.startswith('v'):
714                line = line.replace(era_name, self.name)
715                # Can call simple functions directly.  Those need little modification.
716                if self.name + '(' in line:
717                    line = line.replace(self.name + '(', f"erfa_ufunc.{self.name}(")
718
719            # Call of function that is being tested.
720            elif era_name in line:
721                line = line.replace(era_name, f"erfa_ufunc.{self.name}")
722                # correct for LDBODY (complete hack!)
723                line = line.replace('3, b', 'b').replace('n, b', 'b')
724                # Split into function name and call arguments.
725                start, _, arguments = line.partition('(')
726                # Get arguments, stripping excess spaces and, for numbers, remove
727                # leading zeros since python cannot deal with items like '01', etc.
728                args = []
729                for arg in arguments[:-1].split(','):
730                    arg = arg.strip()
731                    while arg[0] == '0' and len(arg) > 1 and arg[1] in '0123456789':
732                        arg = arg[1:]
733                    args.append(arg)
734                # Get input and output arguments.
735                in_args = [arg.replace('&', '') for arg in args[:self.nin+self.ninout]]
736                out_args = ([arg.replace('&', '') for arg in args[-self.nout-self.ninout:]]
737                            if len(args) > self.nin else [])
738                # If the call assigned something, that will have been the status.
739                # Prepend any arguments assigned in the call.
740                if '=' in start:
741                    line = ', '.join(out_args+[start])
742                else:
743                    line = ', '.join(out_args) + ' = ' + start
744                line = line + '(' + ', '.join(in_args) + ')'
745                if 'astrom' in out_args:
746                    out.append(line)
747                    line = 'astrom = astrom.view(np.recarray)'
748
749            # In some test functions, there are calls to other ERFA functions.
750            # Deal with those in a super hacky way for now.
751            elif line.startswith('eraA'):
752                line = line.replace('eraA', 'erfa_ufunc.a')
753                start, _, arguments = line.partition('(')
754                args = [arg.strip() for arg in arguments[:-1].split(',')]
755                in_args = [arg for arg in args if '&' not in arg]
756                out_args = [arg.replace('&', '') for arg in args if '&' in arg]
757                line = (', '.join(out_args) + ' = '
758                        + start + '(' + ', '.join(in_args) + ')')
759                if 'atioq' in line or 'atio13' in line or 'apio13' in line:
760                    line = line.replace(' =', ', j =')
761
762            # And the same for some other functions, which always have a
763            # 2-element time as inputs.
764            elif line.startswith('eraS'):
765                line = line.replace('eraS', 'erfa_ufunc.s')
766                start, _, arguments = line.partition('(')
767                args = [arg.strip() for arg in arguments[:-1].split(',')]
768                in_args = args[:2]
769                out_args = args[2:]
770                line = (', '.join(out_args) + ' = '
771                        + start + '(' + ', '.join(in_args) + ')')
772
773            # Input number setting.
774            elif '=' in line:
775                # Small clean-up.
776                line = line.replace('=  ', '= ')
777                # Hack to make astrom element assignment work.
778                if line.startswith('astrom'):
779                    out.append('astrom = np.zeros((), erfa_ufunc.dt_eraASTROM).view(np.recarray)')
780                # Change access to p and v elements for double[2][3] pv arrays
781                # that were not caught by the general replacement above (e.g.,
782                # with names not equal to 'pv')
783                name, _, rest = line.partition('[')
784                if (rest and rest[0] in '01' and name in self.var_dtypes
785                        and self.var_dtypes[name] == 'dt_pv'):
786                    line = name + "[" + ("'p'" if rest[0] == "0" else "'v'") + rest[1:]
787
788            out.append(line)
789
790        return out
791
792
793def main(srcdir=DEFAULT_ERFA_LOC, templateloc=DEFAULT_TEMPLATE_LOC, verbose=True):
794    from jinja2 import Environment, FileSystemLoader
795
796    outfn = 'core.py'
797    ufuncfn = 'ufunc.c'
798    testdir = 'tests'
799    testfn = 'test_ufunc.py'
800
801    if verbose:
802        print_ = print
803    else:
804        def print_(*args, **kwargs):
805            return None
806
807    # Prepare the jinja2 templating environment
808    env = Environment(loader=FileSystemLoader(templateloc))
809
810    def prefix(a_list, pre):
811        return [pre+f'{an_element}' for an_element in a_list]
812
813    def postfix(a_list, post):
814        return [f'{an_element}'+post for an_element in a_list]
815
816    def surround(a_list, pre, post):
817        return [pre+f'{an_element}'+post for an_element in a_list]
818    env.filters['prefix'] = prefix
819    env.filters['postfix'] = postfix
820    env.filters['surround'] = surround
821
822    erfa_c_in = env.get_template(ufuncfn + '.templ')
823    erfa_py_in = env.get_template(outfn + '.templ')
824
825    # Prepare the jinja2 test templating environment
826    env2 = Environment(loader=FileSystemLoader(os.path.join(templateloc, testdir)))
827
828    test_py_in = env2.get_template(testfn + '.templ')
829
830    # Extract all the ERFA function names from erfa.h
831    if os.path.isdir(srcdir):
832        erfahfn = os.path.join(srcdir, 'erfa.h')
833        t_erfa_c_fn = os.path.join(srcdir, 't_erfa_c.c')
834        multifilserc = True
835    else:
836        erfahfn = os.path.join(os.path.split(srcdir)[0], 'erfa.h')
837        t_erfa_c_fn = os.path.join(os.path.split(srcdir)[0], 't_erfa_c.c')
838        multifilserc = False
839
840    with open(erfahfn, "r") as f:
841        erfa_h = f.read()
842        print_("read erfa header")
843
844    with open(t_erfa_c_fn, "r") as f:
845        t_erfa_c = f.read()
846        print_("read C tests")
847
848    funcs = OrderedDict()
849    section_subsection_functions = re.findall(
850        r'/\* (\w*)/(\w*) \*/\n(.*?)\n\n', erfa_h,
851        flags=re.DOTALL | re.MULTILINE)
852    for section, subsection, functions in section_subsection_functions:
853        print_(f"{section}.{subsection}")
854
855        if True:
856
857            func_names = re.findall(r' (\w+)\(.*?\);', functions,
858                                    flags=re.DOTALL)
859            for name in func_names:
860                print_(f"{section}.{subsection}.{name}...")
861                if multifilserc:
862                    # easy because it just looks in the file itself
863                    cdir = (srcdir if section != 'Extra' else
864                            templateloc or '.')
865                    funcs[name] = Function(name, cdir)
866                else:
867                    # Have to tell it to look for a declaration matching
868                    # the start of the header declaration, otherwise it
869                    # might find a *call* of the function instead of the
870                    # definition
871                    for line in functions.split(r'\n'):
872                        if name in line:
873                            # [:-1] is to remove trailing semicolon, and
874                            # splitting on '(' is because the header and
875                            # C files don't necessarily have to match
876                            # argument names and line-breaking or
877                            # whitespace
878                            match_line = line[:-1].split('(')[0]
879                            funcs[name] = Function(name, cdir, match_line)
880                            break
881                    else:
882                        raise ValueError("A name for a C file wasn't "
883                                         "found in the string that "
884                                         "spawned it.  This should be "
885                                         "impossible!")
886
887    test_funcs = [TestFunction.from_function(funcs[name], t_erfa_c)
888                  for name in sorted(funcs.keys())]
889
890    funcs = funcs.values()
891
892    # Extract all the ERFA constants from erfam.h
893    erfamhfn = os.path.join(srcdir, 'erfam.h')
894    with open(erfamhfn, 'r') as f:
895        erfa_m_h = f.read()
896    constants = []
897    for chunk in erfa_m_h.split("\n\n"):
898        result = re.findall(r"#define (ERFA_\w+?) (.+?)$", chunk,
899                            flags=re.DOTALL | re.MULTILINE)
900        if result:
901            doc = re.findall(r"/\* (.+?) \*/\n", chunk, flags=re.DOTALL)
902            for (name, value) in result:
903                constants.append(Constant(name, value, doc))
904
905    # TODO: re-enable this when const char* return values and
906    #       non-status code integer rets are possible
907    # #Add in any "extra" functions from erfaextra.h
908    # erfaextrahfn = os.path.join(srcdir, 'erfaextra.h')
909    # with open(erfaextrahfn, 'r') as f:
910    #     for l in f:
911    #         ls = l.strip()
912    #         match = re.match('.* (era.*)\(', ls)
913    #         if match:
914    #             print_("Extra:  {0} ...".format(match.group(1)))
915    #             funcs.append(ExtraFunction(match.group(1), ls, erfaextrahfn))
916
917    print_("Rendering template")
918    erfa_c = erfa_c_in.render(funcs=funcs)
919    erfa_py = erfa_py_in.render(funcs=funcs, constants=constants)
920    test_py = test_py_in.render(test_funcs=test_funcs)
921
922    if outfn is not None:
923        print_(f"Saving to {outfn}, {ufuncfn} and {testfn}")
924        with open(os.path.join(templateloc, outfn), "w") as f:
925            f.write(erfa_py)
926        with open(os.path.join(templateloc, ufuncfn), "w") as f:
927            f.write(erfa_c)
928        with open(os.path.join(templateloc, testdir, testfn), "w") as f:
929            f.write(test_py)
930
931    print_("Done!")
932
933    return erfa_c, erfa_py, funcs, test_py, test_funcs
934
935
936if __name__ == '__main__':
937    from argparse import ArgumentParser
938
939    ap = ArgumentParser()
940    ap.add_argument('srcdir', default=DEFAULT_ERFA_LOC, nargs='?',
941                    help='Directory where the ERFA c and header files '
942                         'can be found or to a single erfa.c file '
943                         '(which must be in the same directory as '
944                         'erfa.h). Default: "{}"'.format(DEFAULT_ERFA_LOC))
945    ap.add_argument('-t', '--template-loc',
946                    default=DEFAULT_TEMPLATE_LOC,
947                    help='the location where the "core.py.templ" and '
948                         '"ufunc.c.templ templates can be found.')
949    ap.add_argument('-q', '--quiet', action='store_false', dest='verbose',
950                    help='Suppress output normally printed to stdout.')
951
952    args = ap.parse_args()
953    main(args.srcdir, args.template_loc, args.verbose)
954