def handle(self, *args, **options):
print "bla for the win"
# get the db, exit if exists
dataset = get_scheme()
if dataset.count() > 0:
if options['clean']:
dataset.delete_many({})
else:
print 'Schema is already uploaded. Exiting.'
return
reader = self._parse_csv(SCHEME_FILENAME)
# Create all the nodes in a linear way.
# We Create dictionary in the following format
# { node_code : [node_object, node_parent_code] }
nodes = {
line['CODE']: [
Tree(name=line['NAME'],
code=line['CODE'],
expense=line['DIRECTION']), line['PARENT']
]
for line in reader
}
root = Tree()
expenditure_root = Tree(name=u'הוצאות',
amount=None,
code=None,
expense=True,
_id=None)
revenue_root = Tree(name=u'הכנסות',
amount=None,
code=None,
expense=False,
_id=None)
root.add_child(expenditure_root)
root.add_child(revenue_root)
# We go through the dict to create the hierarchy.
for [node, parent] in nodes.values():
if not parent:
if node.expense:
expenditure_root.add_child(node)
else:
revenue_root.add_child(node)
else:
nodes[parent][0].add_child(node)
# Upload the Tree to the DB.
dataset.insert(root.to_dict())
dataset.close()
def create_tree(muni, year):
schema_dataset = get_scheme()
tree = Tree.from_dict(schema_dataset.find_one())
tree.update_field('muni', muni)
tree.update_field('year', year)
if tree.children[0].expense:
expense_root = tree.children[0]
revenue_root = tree.children[1]
else:
expense_root = tree.children[1]
revenue_root = tree.children[0]
budget_dataset = get_raw_budget(muni, year)
for i, line in enumerate(budget_dataset.find({})):
node = Tree(muni=muni, year=year, **line)
if len([x for x in expense_root.children if x.code == node.code[1]]):
expense_root.insert_node(node)
else:
revenue_root.insert_node(node)
tree.update_amount()
schema_dataset.close()
budget_dataset.close()
return tree
