<commit_before>#!/usr/local/bin/python2

import sys, os, threading, psycopg2, time
from PyQt5 import QtCore, QtWidgets, QtGui
from subprocess import check_output, call
from mainwindow import Ui_MainWindow

class Shp2Pgsql(QtWidgets.QMainWindow, Ui_MainWindow):
    """docstring for Shp2Pgsql"""
    def __init__(self):
        super(Shp2Pgsql, self).__init__()
        self.setupUi(self)

        path = os.environ['PATH']
        os.environ['PATH'] = '/usr/local/bin:/usr/bin:{}'.format(path)

        self.initUi()
        self.connection = None
        self.counter = 0
        self.file_count = 0
        self.threads = {}
        self.selected_rows = []
        self.home = os.path.expanduser('~')

    def initUi(self):
        self.btn_test.clicked.connect(self.connect)
        self.btn_import.clicked.connect(self.import_)
        self.btn_add.clicked.connect(self.add_file)
        self.btn_clear.clicked.connect(lambda: self.clear_table(1))
        self.btn_clear2.clicked.connect(lambda: self.clear_table(2))
        self.btn_save_log.clicked.connect(self.save_log)
        self.btn_remove.setVisible(False)
        self.btn_remove.clicked.connect(lambda: self.remove_selected(1))
        self.btn_remove2.setVisible(False)
        self.btn_remove2.clicked.connect(lambda: self.remove_selected(2))
        self.btn_fetch.clicked.connect(self.fetch_table)
        self.btn_export.clicked.connect(self.export_)

        self.tbl_table.cellDoubleClicked.connect(self.export_dest)

        self.act_add.triggered.connect(self.add_file)
        self.act_about.triggered.connect(self.show_about)

        self.tbl_file.itemSelectionChanged.connect(lambda: self.selection_changed(1))
        self.tbl_table.itemSelectionChanged.connect(lambda: self.selection_changed(2))

        fixed_font = QtGui.QFontDatabase.systemFont(QtGui.QFontDatabase.FixedFont)
        self.txt_log.setFont(fixed_font)
        self.txt_log.insertPlainText("========== Shp2PgsqlGUI for macOS ==========\n")

        self.show()

    def connect(self, **kwargs):
        host = self.fld_host.text()
        port = self.fld_port.text()
        user = self.fld_user.text()
        password = self.fld_password.text()
        dbname = self.fld_dbname.text()

        output = kwargs.get('output', True)

        try:
            self.connection = psycopg2.connect(host=host, port=port, user=user, password=password, dbname=dbname)
            if output:
                self.alert('Success', 'Connected to database', 'info')
            self.write_log('Connected to database {} on {}:{}'.format(dbname, host, port))
        except Exception as e:
            if output:
                self.alert('Error', 'Could not conect to database', 'critical')

    def add_file(self):
        files, _ = QtWidgets.QFileDialog.getOpenFileNames(self, 'Choose shape file to import', self.home, 'ESRI Shapefiles (*.shp)')
        if len(files):
            row_count = self.tbl_file.rowCount()
            self.tbl_file.setRowCount(len(files) + row_count)

            for i, file in enumerate(files):
                self.tbl_file.setItem(i + row_count, 0, QtWidgets.QTableWidgetItem(file))
                self.tbl_file.setItem(i + row_count, 1, QtWidgets.QTableWidgetItem('public'))
                self.tbl_file.setItem(i + row_count, 2, QtWidgets.QTableWidgetItem(os.path.basename(os.path.splitext(file)[0]).lower()))
                self.tbl_file.setItem(i + row_count, 3, QtWidgets.QTableWidgetItem('geom'))
                self.tbl_file.setItem(i + row_count, 4, QtWidgets.QTableWidgetItem('4326'))

            self.tbl_file.resizeColumnsToContents()

    def clear_table(self, num):
        if num == 1:
            table = self.tbl_file
        else:
            table = self.tbl_table

        table.clearContents()
        table.setRowCount(0)
        table.resizeColumnsToContents()

    def selection_changed(self, num):
        items = self.tbl_file.selectedItems() if num == 1 else self.tbl_table.selectedItems()
        self.selected_rows = []
        for item in items:
            row = item.row()
            if not row in self.selected_rows:
                self.selected_rows.append(row)

        if num == 1:
            self.btn_remove.setVisible(len(self.selected_rows) > 0)
        else:
            self.btn_remove2.setVisible(len(self.selected_rows) > 0)

    def remove_selected(self, num):
        self.selected_rows.sort(reverse=True)
        for row in self.selected_rows:
            if num == 1:
                self.tbl_file.removeRow(row)
            else:
                self.tbl_table.removeRow(row)
        self.selected_rows = []

    def fetch_table(self):
        try:
            self.connect(output=False)
            if not self.connection:
                self.alert('Error', 'No database connection', 'critical')
                return
        except Exception:
            self.alert('Error', 'No database connection', 'critical')
            return

        cur = self.connection.cursor()
        cur.execute("select table_schema, table_name from information_schema.tables where table_name != 'spatial_ref_sys' and table_schema not in ('pg_catalog', 'information_schema') and table_type = 'BASE TABLE'")
        tables = cur.fetchall()

        if tables:
            self.tbl_table.setRowCount(len(tables))
            for i, t in enumerate(tables):
                self.tbl_table.setItem(i, 0, QtWidgets.QTableWidgetItem(t[0]))
                self.tbl_table.setItem(i, 1, QtWidgets.QTableWidgetItem(t[1]))
                self.tbl_table.setItem(i, 2, QtWidgets.QTableWidgetItem('geom'))
                self.tbl_table.setItem(i, 3, QtWidgets.QTableWidgetItem(os.path.join(self.home, t[1] + '.shp')))
        else:
            self.write_log('No table found')

        self.tbl_table.resizeColumnsToContents()
        cur.close()
        self.connection.close()
        self.connection = None

    def export_dest(self, row, col):
        if col == 3:
            file, _ = QtWidgets.QFileDialog.getSaveFileName(self, 'Save As', self.home, 'ESRI Shapefile (*.shp)')<commit_msg>automatic check geography or geometry column and only add tables containing it to export list<commit_after>#!/usr/local/bin/python2

import sys, os, threading, psycopg2, time
from PyQt5 import QtCore, QtWidgets, QtGui
from subprocess import check_output, call
from mainwindow import Ui_MainWindow

class Shp2Pgsql(QtWidgets.QMainWindow, Ui_MainWindow):
    """docstring for Shp2Pgsql"""
    def __init__(self):
        super(Shp2Pgsql, self).__init__()
        self.setupUi(self)

        path = os.environ['PATH']
        os.environ['PATH'] = '/usr/local/bin:/usr/bin:{}'.format(path)

        self.initUi()
        self.connection = None
        self.counter = 0
        self.file_count = 0
        self.threads = {}
        self.selected_rows = []
        self.home = os.path.expanduser('~')

    def initUi(self):
        self.btn_test.clicked.connect(self.connect)
        self.btn_import.clicked.connect(self.import_)
        self.btn_add.clicked.connect(self.add_file)
        self.btn_clear.clicked.connect(lambda: self.clear_table(1))
        self.btn_clear2.clicked.connect(lambda: self.clear_table(2))
        self.btn_save_log.clicked.connect(self.save_log)
        self.btn_remove.setVisible(False)
        self.btn_remove.clicked.connect(lambda: self.remove_selected(1))
        self.btn_remove2.setVisible(False)
        self.btn_remove2.clicked.connect(lambda: self.remove_selected(2))
        self.btn_fetch.clicked.connect(self.fetch_table)
        self.btn_export.clicked.connect(self.export_)

        self.tbl_table.cellDoubleClicked.connect(self.export_dest)

        self.act_add.triggered.connect(self.add_file)
        self.act_about.triggered.connect(self.show_about)

        self.tbl_file.itemSelectionChanged.connect(lambda: self.selection_changed(1))
        self.tbl_table.itemSelectionChanged.connect(lambda: self.selection_changed(2))

        fixed_font = QtGui.QFontDatabase.systemFont(QtGui.QFontDatabase.FixedFont)
        self.txt_log.setFont(fixed_font)
        self.txt_log.insertPlainText("========== Shp2PgsqlGUI for macOS ==========\n")

        self.show()

    def connect(self, **kwargs):
        host = self.fld_host.text()
        port = self.fld_port.text()
        user = self.fld_user.text()
        password = self.fld_password.text()
        dbname = self.fld_dbname.text()

        output = kwargs.get('output', True)

        try:
            self.connection = psycopg2.connect(host=host, port=port, user=user, password=password, dbname=dbname)
            if output:
                self.alert('Success', 'Connected to database', 'info')
            self.write_log('Connected to database {} on {}:{}'.format(dbname, host, port))
        except Exception as e:
            if output:
                self.alert('Error', 'Could not conect to database', 'critical')

    def add_file(self):
        files, _ = QtWidgets.QFileDialog.getOpenFileNames(self, 'Choose shape file to import', self.home, 'ESRI Shapefiles (*.shp)')
        if len(files):
            row_count = self.tbl_file.rowCount()
            self.tbl_file.setRowCount(len(files) + row_count)

            for i, file in enumerate(files):
                self.tbl_file.setItem(i + row_count, 0, QtWidgets.QTableWidgetItem(file))
                self.tbl_file.setItem(i + row_count, 1, QtWidgets.QTableWidgetItem('public'))
                self.tbl_file.setItem(i + row_count, 2, QtWidgets.QTableWidgetItem(os.path.basename(os.path.splitext(file)[0]).lower()))
                self.tbl_file.setItem(i + row_count, 3, QtWidgets.QTableWidgetItem('geom'))
                self.tbl_file.setItem(i + row_count, 4, QtWidgets.QTableWidgetItem('4326'))

            self.tbl_file.resizeColumnsToContents()

    def clear_table(self, num):
        if num == 1:
            table = self.tbl_file
        else:
            table = self.tbl_table

        table.clearContents()
        table.setRowCount(0)
        table.resizeColumnsToContents()

    def selection_changed(self, num):
        items = self.tbl_file.selectedItems() if num == 1 else self.tbl_table.selectedItems()
        self.selected_rows = []
        for item in items:
            row = item.row()
            if not row in self.selected_rows:
                self.selected_rows.append(row)

        if num == 1:
            self.btn_remove.setVisible(len(self.selected_rows) > 0)
        else:
            self.btn_remove2.setVisible(len(self.selected_rows) > 0)

    def remove_selected(self, num):
        self.selected_rows.sort(reverse=True)
        for row in self.selected_rows:
            if num == 1:
                self.tbl_file.removeRow(row)
            else:
                self.tbl_table.removeRow(row)
        self.selected_rows = []

    def fetch_table(self):
        try:
            self.connect(output=False)
            if not self.connection:
                self.alert('Error', 'No database connection', 'critical')
                return
        except Exception:
            self.alert('Error', 'No database connection', 'critical')
            return

        cur = self.connection.cursor()
        cur.execute("select table_schema, table_name from information_schema.tables where table_name != 'spatial_ref_sys' and table_schema not in ('pg_catalog', 'information_schema') and table_type = 'BASE TABLE'")
        tables = cur.fetchall()

        if tables:
            self.tbl_table.setRowCount(0)
            for t in tables:
                cur.execute("select column_name, udt_name from information_schema.columns where table_name='{}' and udt_name in ('geography', 'geometry')".format(t[1]))
                geom = cur.fetchone()
                if geom:
                    row_count = self.tbl_table.rowCount()
                    self.tbl_table.insertRow(row_count)
                    self.tbl_table.setItem(row_count, 0, QtWidgets.QTableWidgetItem(t[0]))
                    self.tbl_table.setItem(row_count, 1, QtWidgets.QTableWidgetItem(t[1]))
                    self.tbl_table.setItem(row_count, 2, QtWidgets.QTableWidgetItem(geom[0]))
                    self.tbl_table.setItem(row_count, 3, QtWidgets.QTableWidgetItem(os.path.join(self.home, t[1] + '.shp')))
                    self.write_log('{} column found in table {}.{}. Adding table to list.'.format(geom[1], t[0], t[1]))
        else:
            self.write_log('No table found')

        if not self.tbl_table.rowCount():
            self.write_log('Could not find table with geography or geometry column.')
            self.alert('Info', 'Could not find table with geography or geometry column.', 'info')

        self.tbl_table.resizeColumnsToContents()
        cur.close()
        self.connection.close()
        self.connection = None

    def export_dest(self, row, col):
        if col == 3:
            file, _ = QtWidgets.QFileDialog.getSaveFileName(self, 'Save As', self.home, 'ESRI Shapefile (*.shp)')
{'fraction_non_alphanumeric': 0.1186845348120764, 'fraction_numerical': 0.0073937153419593345, 'mean_word_length': 2.498114224137931, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6562369', 'n_tokens_mistral': 4098, 'n_tokens_neox': 3943, 'n_words': 717}