Initial commit

3 years ago · 4fd2698947
4 changed files with 389 additions and 0 deletions
--- a/.gitignore
+++ b/.gitignore
@ -0,0 +1,12 @@
+# Python
+__pycache__
+
+# Data
+datafiles/*
+!datafiles/mapping.example.json
+
+# Output
+trains.sqlite
+
+# Misc
+view-train.sql
--- a/README.md
+++ b/README.md
@ -0,0 +1,101 @@
+# CFR XML to SQL
+
+Un script Python pentru convertirea
+[fișierelor XML oferite de S.C. Informatica Feroviară S.A.](https://data.gov.ro/dataset?organization=sc-informatica-feroviara-sa)
+într-o bază de date SQLite.
+
+## Cum se folosește
+
+- se plasează fișierele `.xml` în folderul `datafiles`
+- (opțional) se completează fișierul `datafiles/mapping.json` cu numele companiilor
+- se rulează `convert.py`:
+
+```bash
+# Unix
+./convert.py
+```
+
+```powershell
+# Windows
+python convert.py
+```
+
+Scriptul este scris pentru Python 3.
+
+## Schemă bază de date
+
+Schema este validă pentru versiunea 2.
+
+Coloanele fără tip de date sunt luate direct din XML (deci cel mai probabil tip text).
+
+### Tabel **Meta**
+
+Coloană  | Tip de date | Descriere
+---------|-------------|----------
+Versiune | int         | Versiunea bazei de date (schema aceasta descrie versiunea 2)
+
+### Tabel **Companii**
+
+Coloană   | Tip de date | Descriere
+----------|-------------|----------
+Id        | int pk      |
+NumeLegal | text        | Numele oficial al companiei *(ex: S.C. Companie S.R.L)*
+NumeComun | text        | Numele după care este cunoscută compania *(în general numele oficial fără prefix sau sufix; ex: Companie)*
+
+### Tabel **Trenuri**
+
+Coloană       | Tip de date | Descriere
+--------------|-------------|----------
+Number        | int pk      |
+IdCompanie    | int         | fk Companii->Id
+CategorieTren |             | R/IR/etc.
+KmCum         | int         | Lungimea totală a traseului în **metri**
+Lungime       | int         | Lungimea trenului
+Numar         |             |
+Operator      |             | ? - Uneori conține un număr unic pentru un anumit operator, alteori este gol
+Proprietar    |             |
+Putere        |             |
+Rang          |             |
+Servicii      |             |
+Tonaj         |             |
+
+### Tabel **Trase**
+
+Coloană           | Tip de date | Descriere
+------------------|-------------|----------
+NumarTren         | int         | fk Trenuri->Number
+Id                | int         | Id trasă (1, 2, ...)
+Tip               |             |
+CodStatieInitiala | int         | fk Statii->CodStatie
+CodStatieFinala   | int         | fk Statii->CodStatie
+
+### Tabel **Statii**
+
+Coloană   | Tip de date | Descriere
+----------|-------------|----------
+CodStatie | int pk      | Codul folosit în setul de date
+Denumire  | text        | Numele stației *(ex: București Nord Gr.A)*
+
+### Tabel **ElementeTrasa**
+
+Coloană           | Tip de date | Descriere
+------------------|-------------|----------
+NumarTren         | int         | fk Trenuri->Number + Trase->NumarTren
+IdTrasa           | int         | fk Trase->Id
+Secventa          | int         | Ordinea secvenței în trasă (1, 2, ...)
+Ajustari          |             |
+CodStaDest        |             | fk Statii->CodStatie
+CodStaOrigine     |             | fk Statii->CodStatie
+DenStaDestinatie  |             |
+DenStaOrigine     |             |
+Km                | int         | Distanța în **metri** între stații *(ex: 3022 -> 3022 m / 3 km)*
+Lungime           | int         | Lungimea trenului
+OraP              | int         | Ora plecării (numărul de secunde de la 00:00:00 în ziua primei plecări)
+OraS              | int         | Ora sosirii (numărul de secunde de la 00:00:00 în ziua primei plecări)
+Rci               |             |
+Rco               |             |
+Restrictie        |             |
+StationareSecunde | int         | Numărul de secunde de staționare *în stația origine*
+TipOprire         |             |
+Tonaj             |             |
+VitezaLivret      | int         | Viteza între stația origine și destinație în km/h
--- a/convert.py
+++ b/convert.py
@ -0,0 +1,266 @@
+#! /usr/bin/env python3
+
+import xml.etree.ElementTree as ET
+from itertools import takewhile
+try:
+	from tqdm import tqdm
+except ImportError:
+	def tqdm(iter, *args, **kargs):
+		return iter
+
+def get_database_connection():
+	import sqlite3
+	return sqlite3.connect('trains.sqlite')
+
+def create_tables(con):
+	cursor = con.cursor()
+	
+	cursor.execute("select name from sqlite_master where type='table' order by name;")
+	tables = [item[0] for item in cursor.fetchall()]
+
+	if 'Meta' not in tables:
+		cursor.execute('create table Meta (Versiune int)')
+	if 'Companii' not in tables:
+		cursor.execute('create table Companii (Id integer primary key, NumeLegal text, NumeComun text)')
+	if 'Trenuri' not in tables:
+		cursor.execute('create table Trenuri (Number integer primary key, IdCompanie int, CategorieTren, KmCum int, Lungime int, Numar, Operator, Proprietar, Putere, Rang, Servicii, Tonaj)')
+	if 'Trase' not in tables:
+		cursor.execute('create table Trase (NumarTren int, Id int, Tip, CodStatieInitiala int, CodStatieFinala int)')
+	if 'ElementeTrasa' not in tables:
+		cursor.execute('create table ElementeTrasa (NumarTren int, IdTrasa int, Secventa int, Ajustari, CodStaDest, CodStaOrigine, DenStaDestinatie, DenStaOrigine, Km int, Lungime int, OraP int, OraS int, Rci, Rco, Restrictie, StationareSecunde int, TipOprire, Tonaj, VitezaLivret int)')
+	if 'Statii' not in tables:
+		cursor.execute('create table Statii (CodStatie integer primary key, Denumire text)')
+	con.commit()
+
+def insert(con, table, *args, _commit=True, **kargs):
+	cursor = con.cursor()
+
+	if args and not kargs:
+		arg_str = '(' + ', '.join((['?'] * len(args))) + ')'
+		cursor.execute(f"insert into {table} values {arg_str}", args)
+	elif not args and kargs:
+		arg_str = '(' + ', '.join((['?'] * len(kargs))) + ')'
+		apair = list(kargs.items())
+		keys = [k for (k, _) in apair]
+		values = [v for (_, v) in apair]
+		columns = '(' + ','.join(keys) + ')'
+		cursor.execute(f"insert into {table} {columns} values {arg_str}", values)
+	else:
+		raise Exception('Provide args XOR kargs')
+
+	if _commit:
+		con.commit()
+
+def get_data_folder():
+	import os
+	data_folder = os.path.join('.', 'datafiles')
+	return data_folder
+
+def get_xml_files():
+	import os
+	data_folder = get_data_folder()
+	for entry in os.listdir(data_folder):
+		entry = os.path.join(data_folder, entry)
+		if os.path.isfile(entry):
+			if os.path.splitext(entry)[1] == '.xml':
+				yield entry
+
+def get_mappings():
+	from os.path import join
+	import json
+	data_folder = get_data_folder()
+	mappings_file = join(data_folder, 'mapping.json')
+	try:
+		with open(mappings_file) as f:
+			fj = json.load(f)
+			return fj['mappings']
+	except:
+		return []
+			
+def train_number_stoi(s):
+	return int(''.join(takewhile(lambda c: c.isnumeric(), s)))
+
+def find_trains(con):
+	cursor = con.cursor()
+
+	cursor.execute('select Number from Trenuri;')
+
+	return set((item[0] for item in cursor.fetchall()))
+
+def find_trase(con, train_number=None):
+	cursor = con.cursor()
+
+	if train_number is None:
+		cursor.execute('select NumarTren, Id from Trase')
+	else:
+		cursor.execute('select NumarTren, Id from Trase where NumarTren = ?', (train_number,))
+
+	return [(nt, i) for nt, i in cursor.fetchall()]
+
+def find_secvente(con, train_number=None, id_trasa=None):
+	cursor = con.cursor()
+
+	if train_number is None:
+		cursor.execute('select NumarTren, IdTrasa, Secventa from ElementeTrasa')
+	elif id_trasa is None:
+		cursor.execute('select NumarTren, IdTrasa, Secventa from ElementeTrasa where NumarTren = ?', (train_number,))
+	else:
+		cursor.execute('select NumarTren, IdTrasa, Secventa from ElementeTrasa where NumarTren = ? and IdTrasa = ?', (train_number, id_trasa))
+
+	return [(nt, it, s) for nt, it, s in cursor.fetchall()]
+
+def find_station_ids(con):
+	cursor = con.cursor()
+
+	cursor.execute('select CodStatie from Statii;')
+
+	return set((item[0] for item in cursor.fetchall()))
+
+def find_companies(con):
+	cursor = con.cursor()
+	
+	cursor.execute('select Id, NumeLegal, NumeComun from Companii')
+
+	return list(cursor.fetchall())
+		
+
+def main():
+	con = get_database_connection()
+	create_tables(con)
+	insert(con, 'Meta', 2)
+
+	station_ids = find_station_ids(con)
+	companies = find_companies(con)
+
+	mappings = get_mappings()
+
+	def get_company_name(path):
+		try:
+			from os.path import basename
+			bn = basename(path)
+			for mapping in mappings:
+				if mapping['filename'] == bn:
+					return mapping['legalName'], mapping['commonName']
+		except:
+			pass
+		return None, None
+
+
+	for f in get_xml_files():
+		company_legal_name, company_common_name = get_company_name(f)
+		if len([cln for (_, cln, ccn) in companies if cln == company_legal_name and ccn == company_common_name]) == 0:
+			insert(con, 'Companii', NumeLegal=company_legal_name, NumeComun=company_common_name)
+			companies = find_companies(con)
+		company_id = [i for (i, cln, ccn) in companies if cln == company_legal_name and ccn == company_common_name][0]
+
+		tree = ET.parse(f)
+		el_trenuri = tree.find("/XmlMts/Mt/Trenuri")
+		
+		trains = find_trains(con)
+
+		print(f'Adding {company_common_name or f}...')
+		for el_tren in tqdm(el_trenuri.findall("./Tren")):
+			train_number_str = el_tren.attrib['Numar']
+			train_number = train_number_stoi(train_number_str)
+
+			if train_number in trains:
+				continue
+
+			trase = find_trase(con, train_number)
+
+			insert(
+				con, 
+				'Trenuri',
+				train_number,
+				company_id,
+				el_tren.attrib['CategorieTren'],
+				el_tren.attrib['KmCum'],
+				el_tren.attrib['Lungime'],
+				train_number_str,
+				el_tren.attrib['Operator'],
+				el_tren.attrib['Proprietar'],
+				el_tren.attrib['Putere'],
+				el_tren.attrib['Rang'],
+				el_tren.attrib['Servicii'],
+				el_tren.attrib['Tonaj'],
+				_commit=False,
+			)
+
+			for el_trasa in el_tren.findall('./Trase/Trasa'):
+				id_trasa = int(el_trasa.attrib['Id'])
+
+				if (train_number, id_trasa) in trase:
+					continue
+
+				secvente = find_secvente(con, train_number, id_trasa)
+
+				insert(
+					con,
+					'Trase',
+					train_number,
+					id_trasa,
+					el_trasa.attrib['Tip'],
+					el_trasa.attrib['CodStatieInitiala'],
+					el_trasa.attrib['CodStatieFinala'],
+					_commit=False,
+				)
+
+				for el_elementtrasa in el_trasa.findall('./ElementTrasa'):
+					secventa = int(el_elementtrasa.attrib['Secventa'])
+
+					if (train_number, id_trasa, secventa) in secvente:
+						continue
+				
+					insert(
+						con,
+						'ElementeTrasa',
+						train_number,
+						id_trasa,
+						secventa,
+						el_elementtrasa.attrib['Ajustari'],
+						el_elementtrasa.attrib['CodStaDest'],
+						el_elementtrasa.attrib['CodStaOrigine'],
+						el_elementtrasa.attrib['DenStaDestinatie'],
+						el_elementtrasa.attrib['DenStaOrigine'],
+						el_elementtrasa.attrib['Km'],
+						el_elementtrasa.attrib['Lungime'],
+						el_elementtrasa.attrib['OraP'],
+						el_elementtrasa.attrib['OraS'],
+						el_elementtrasa.attrib['Rci'],
+						el_elementtrasa.attrib['Rco'],
+						el_elementtrasa.attrib['Restrictie'],
+						el_elementtrasa.attrib['StationareSecunde'],
+						el_elementtrasa.attrib['TipOprire'],
+						el_elementtrasa.attrib['Tonaj'],
+						el_elementtrasa.attrib['VitezaLivret'],
+						_commit=False,
+					)
+
+					if el_elementtrasa.attrib['CodStaOrigine'].isnumeric():
+						cod_sta_orig = int(el_elementtrasa.attrib['CodStaOrigine'])
+						if cod_sta_orig not in station_ids:
+							station_ids.add(cod_sta_orig)
+							insert(
+								con,
+								'Statii',
+								cod_sta_orig,
+								el_elementtrasa.attrib['DenStaOrigine'],
+								_commit=False,
+							)
+					if el_elementtrasa.attrib['CodStaDest'].isnumeric():
+						cod_sta_orig = int(el_elementtrasa.attrib['CodStaDest'])
+						if cod_sta_orig not in station_ids:
+							station_ids.add(cod_sta_orig)
+							insert(
+								con,
+								'Statii',
+								cod_sta_orig,
+								el_elementtrasa.attrib['DenStaDestinatie'],
+								_commit=False,
+							)
+		con.commit()
+
+	con.commit()
+
+if __name__ == '__main__':
+	main()
--- a/datafiles/mapping.example.json
+++ b/datafiles/mapping.example.json
@ -0,0 +1,10 @@
+{
+	"$schema": "https://cfr-xml-to-sql.extras.dcdevelop.xyz/mapping.schema.json",
+	"mappings": [
+		{
+			"filename": "example.xml",
+			"legalName": "S.C. Example Company S.R.L",
+			"commonName": "Example Company"
+		}
+	]
+}