Forms.py

#!/usr/bin/env python3

"""
Module for creating forms in a SpecDB database
"""

import io
from io import StringIO
import json
import os
import re

import ruamel.yaml

def collect_schema_comments(tables=None, num=None):
	"""
	Collect the schema and comments for requested tables in SpecDB schema
	
	Parameters
	----------
	+ tables	list, table names to collect schema and comments for
	
	Returns
	-------
	+ form_dic	ruamel.yaml dictionary with comments added
	"""
	
	if tables is None:
		print("Argument `table` cannot be none")
		print("Aborting")
		sys.exit()
	
	if num is not None:
		if len(tables) != len(num):
			print("the num argument must be same length as table")
			print("Aborting")
			sys.exit()
	
	specdb_path = os.path.abspath(os.path.dirname(__file__))
	sql_path    = os.path.join(specdb_path, '../sql/specdb.sql')
	
	form_dic   = ruamel.yaml.comments.CommentedMap()
	table_name = ""
	
	with open(sql_path, 'r') as fp:
		
		for line in fp.readlines():
			line = line.rstrip()
			
			if re.search(r'^CREATE ', line):
				tr = line.split()
				table_name = tr[2]
				
				if table_name in tables:
				
					comment_split = line.split(' -- ')
					comment = comment_split[-1]
					form_dic[table_name] = ruamel.yaml.comments.CommentedMap()
					form_dic.yaml_add_eol_comment(
						comment,
						table_name,
						column=25)
					#continue
				
			elif re.search(r' -- ', line):
			
				if table_name in tables:
					sql_col  = line.split(' -- ')
					col_info = sql_col[0].split()
					comment  = sql_col[-1]
					
					column   = col_info[0]
					sql_type = col_info[1]
					
					form_dic[table_name][column] = None
					form_dic[table_name].yaml_add_eol_comment(
						comment, column, column=35)
	
	return form_dic 

def forms(table=None, num=None, input_dict=None):
	"""
	Create a YAML form for SpecDB tables
	
	Parameters
	----------
	+ table		list, table names to generate forms for
	+ num		list, number of forms to make for each table
				default 1, if any number is given then len(table) == len(num)
	+ input_dic	dictionary to add comments to -- not required
	
	Returns
	-------
	form_out	String of YAML form
	"""
	
	if table is None:
		print("Argument `table` cannot be none")
		print("Aborting")
		sys.exit()
	
	if num is not None:
		if len(table) != len(num):
			print("the num argument must be same length as table")
			print("Aborting")
			sys.exit()
	
	#print(table)
	form_dic = collect_schema_comments(tables=table, num=num)
	#print(json.dumps(form_dic,indent=2))
	yaml = ruamel.yaml.YAML()
	yaml.preserve_quotes = True
	
	if num is not None: num = [int(i) for i in num]
	else: 			    num = [1 for i in range(len(table))]
	
	form = ruamel.yaml.comments.CommentedMap()
	for i, tbl in enumerate(table):
		if tbl not in form_dic:
			if input_dict:
				form[tbl] = input_dict[tbl]
			continue
		form[tbl] = ruamel.yaml.comments.CommentedMap()
		#print(tbl)
		
		comment = form_dic.ca.items[tbl][2].value
		form.yaml_add_eol_comment(comment, tbl, column=25)
		
		for x in range(num[i]):
			form[tbl][x] = ruamel.yaml.comments.CommentedMap()
			for k,v in form_dic[tbl].ca.items.items():
				if input_dict:
					print(tbl,x,k)
					if k not in input_dict[tbl][str(x)]: continue
				comment = v[2].value[1:]
				if input_dict: form[tbl][x][k] = input_dict[tbl][str(x)][k]
				else:          form[tbl][x][k] = form_dic[tbl][k]
				form[tbl][x].yaml_add_eol_comment(comment, k, column=50)
	
	string_stream = StringIO()
	yaml.dump(form, string_stream)
	form_out = string_stream.getvalue()
	string_stream.close()
	return form_out