<commit_before>import os
import subprocess
import re
import codecs
import shutil

import sublime

ST3 = int(sublime.version()) >= 3006

try:
    from . import git_helper
    from .view_collection import ViewCollection
except (ImportError, ValueError):
    import git_helper
    from view_collection import ViewCollection

class GitGutterHandler:
    git_binary_path_error_shown = False

    def __init__(self, view):
        self.load_settings()
        self.view = view
        self.git_temp_file = ViewCollection.git_tmp_file(self.view)
        self.buf_temp_file = ViewCollection.buf_tmp_file(self.view)
        self.git_tree = None
        self.git_dir = None
        self.git_path = None

    def _get_view_encoding(self):
        # get encoding and clean it for python ex: "Western (ISO 8859-1)"
        # NOTE(maelnor): are we need regex here?
        pattern = re.compile(r'.+\((.*)\)')
        encoding = self.view.encoding()
        if encoding == "Undefined":
            encoding = self.view.settings().get('default_encoding')
        if pattern.match(encoding):
            encoding = pattern.sub(r'\1', encoding)

        encoding = encoding.replace('with BOM', '')
        encoding = encoding.replace('Windows', 'cp')
        encoding = encoding.replace('-', '_')
        encoding = encoding.replace(' ', '')

        # work around with ConvertToUTF8 plugin
        origin_encoding = self.view.settings().get('origin_encoding')
        return origin_encoding or encoding

    def on_disk(self):
        # if the view is saved to disk
        on_disk = self.view.file_name() is not None
        if on_disk:
            self.git_tree = self.git_tree or git_helper.git_tree(self.view)
            self.git_dir = self.git_dir or git_helper.git_dir(self.git_tree)
            self.git_path = self.git_path or git_helper.git_file_path(
                self.view, self.git_tree
            )
        return on_disk

    def reset(self):
        if self.on_disk() and self.git_path and self.view.window():
            self.view.window().run_command('git_gutter')

    def get_git_path(self):
        return self.git_path

    def update_buf_file(self):
        chars = self.view.size()
        region = sublime.Region(0, chars)

        # Try conversion
        try:
            contents = self.view.substr(
                region).encode(self._get_view_encoding())
        except UnicodeError:
            # Fallback to utf8-encoding
            contents = self.view.substr(region).encode('utf-8')
        except LookupError:
            # May encounter an encoding we don't have a codec for
            contents = self.view.substr(region).encode('utf-8')

        contents = contents.replace(b'\r\n', b'\n')
        contents = contents.replace(b'\r', b'\n')

        with open(self.buf_temp_file, 'wb') as f:
            if self.view.encoding() == "UTF-8 with BOM":
                f.write(codecs.BOM_UTF8)

            f.write(contents)

    def update_git_file(self):
        # the git repo won't change that often
        # so we can easily wait 5 seconds
        # between updates for performance
        if ViewCollection.git_time(self.view) > 5:
            with open(self.git_temp_file, 'w'):
                pass

            args = [
                self.git_binary_path,
                '--git-dir=' + self.git_dir,
                '--work-tree=' + self.git_tree,
                'show',
                ViewCollection.get_compare(self.view) + ':' + self.git_path,
            ]
            try:
                contents = self.run_command(args)
                contents = contents.replace(b'\r\n', b'\n')
                contents = contents.replace(b'\r', b'\n')
                with open(self.git_temp_file, 'wb') as f:
                    f.write(contents)

                ViewCollection.update_git_time(self.view)
            except Exception:
                pass

    def total_lines(self):
        chars = self.view.size()
        region = sublime.Region(0, chars)
        lines = self.view.lines(region)
        return len(lines)

    # Parse unified diff with 0 lines of context.
    # Hunk range info format:
    #   @@ -3,2 +4,0 @@
    #     Hunk originally starting at line 3, and occupying 2 lines, now
    #     starts at line 4, and occupies 0 lines, i.e. it was deleted.
    #   @@ -9 +10,2 @@
    #     Hunk size can be omitted, and defaults to one line.
    # Dealing with ambiguous hunks:
    #   "A\nB\n" -> "C\n"
    #   Was 'A' modified, and 'B' deleted? Or 'B' modified, 'A' deleted?
    #   Or both deleted? To minimize confusion, let's simply mark the
    #   hunk as modified.
    def process_diff(self, diff_str):
        inserted = []
        modified = []
        deleted = []
        hunk_re = '^@@ \-(\d+),?(\d*) \+(\d+),?(\d*) @@'
        hunks = re.finditer(hunk_re, diff_str, re.MULTILINE)
        for hunk in hunks:
            start = int(hunk.group(3))
            old_size = int(hunk.group(2) or 1)
            new_size = int(hunk.group(4) or 1)
            if not old_size:
                inserted += range(start, start + new_size)
            elif not new_size:
                deleted += [start + 1]
            else:
                modified += range(start, start + new_size)
        return (inserted, modified, deleted)

    def diff(self):
        if self.on_disk() and self.git_path:
            self.update_git_file()
            self.update_buf_file()
            args = [
                self.git_binary_path, 'diff', '-U0', '--no-color', '--no-index',
                self.ignore_whitespace,
                self.patience_switch,
                self.git_temp_file,
                self.buf_temp_file,
            ]
            args = list(filter(None, args))  # Remove empty args
            results = self.run_command(args)
            encoding = self._get_view_encoding()
            try:
                decoded_results = results.decode(encoding.replace(' ', ''))
            except UnicodeError:
                try:
                    decoded_results = results.decode("utf-8")
                except UnicodeDecodeError:
                    decoded_results = ""
            except LookupError:
                try:
                    decoded_results = codecs.decode(results)
                except UnicodeDecodeError:
                    decoded_results = ""
            return self.process_diff(decoded_results)
        else:
            return ([], [], [])

    def untracked(self):
        return self.handle_files([])

    def ignored(self):
        return self.handle_files(['-i'])

    def handle_files(self, additionnal_args):
        if self.on_disk() and self.git_path:
            args = [
                self.git_binary_path,
                '--git-dir=' + self.git_dir,
                '--work-tree=' + self.git_tree,
                'ls-files', '--other', '--exclude-standard',
            ] + additionnal_args + [
                os.path.join(self.git_tree, self.git_path),
            ]
            args = list(filter(None, args))  # Remove empty args
            results = self.run_command(args)
            encoding = self._get_view_encoding()
            try:
                decoded_results = results.decode(encoding.replace(' ', ''))
            except UnicodeError:
                decoded_results = results.decode("utf-8")
            return (decoded_results != "")
        else:
            return False

    def git_commits(self):
        args = [
            self.git_binary_path,
            '--git-dir=' + self.git_dir,
            '--work-tree=' + self.git_tree,
            'log', '--all',
            '--pretty=%s\a%h %an <%aE>\a%ad (%ar)',
            '--date=local', '--max-count=9000'
        ]
        results = self.run_command(args)
        return results

    def git_branches(self):
        args = [
            self.git_binary_path,
            '--git-dir=' + self.git_dir,
            '--work-tree=' + self.git_tree,
            'for-each-ref',
            '--sort=-committerdate',
            '--format=%(subject)\a%(refname)\a%(objectname)',
            'refs/heads/'
        ]
        results = self.run_command(args)
        return results

    def git_tags(self):
        args = [
            self.git_binary_path,
            '--git-dir=' + self.git_dir,
            '--work-tree=' + self.git_tree,
            'show-ref',
            '--tags',
            '--abbrev=7'
        ]
        results = self.run_command(args)
        return results

    def git_current_branch(self):
        args = [
            self.git_binary_path,
            '--git-dir=' + self.git_dir,
            '--work-tree=' + self.git_tree,
            'rev-parse',
            '--abbrev-ref',
            'HEAD'
        ]
        result = self.run_command(args)
        return result

    def run_command(self, args):
        startupinfo = None
        if os.name == 'nt':
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        proc = subprocess.Popen(args, stdout=subprocess.PIPE,
                                startupinfo=startupinfo, stderr=subprocess.PIPE)
        return proc.stdout.read()

    def load_settings(self):
        self.settings = sublime.load_settings('GitGutter.sublime-settings')
        self.user_settings = sublime.load_settings(
            'Preferences.sublime-settings')

        # Git Binary Setting
        git_binary_setting = self.user_settings.get("git_binary") or \
                             self.settings.get("git_binary")
        if isinstance(git_binary_setting, dict):
            self.git_binary_path = git_binary_setting.get(sublime.platform())
            if not self.git_binary_path:
                self.git_binary_path = git_binary_setting.get('default')
        else:
            self.git_binary_path = git_binary_setting

        if self.git_binary_path:
            self.git_binary_path = os.path.expandvars(self.git_binary_path)
        elif ST3:
            self.git_binary_path = shutil.which("git")

        if not self.git_binary_path:
            if not GitGutterHandler.git_binary_path_error_shown:
                GitGutterHandler.git_binary_path_error_shown = True
                msg = ("Your Git binary cannot be found.  If it is installed, add it "
                       "to your PATH environment variable, or add a `git_binary` setting "
                       "in the `User/GitGutter.sublime-settings` file.")
                sublime.error_message(msg)
                raise ValueError("Git binary not found.")

        # Ignore White Space Setting
        self.ignore_whitespace = self.settings.get('ignore_whitespace')
        if self.ignore_whitespace == 'all':
            self.ignore_whitespace = '-w'
        elif self.ignore_whitespace == 'eol':
            self.ignore_whitespace = '--ignore-space-at-eol'
        else:
            self.ignore_whitespace = ''

        # Patience Setting
        self.patience_switch = ''
        patience = self.settings.get('patience')
        if patience:
            self.patience_switch = '--patience'

        # Untracked files
        self.show_untracked = self.settings.get(
            'show_markers_on_untracked_file')

        # Show in minimap
        self.show_in_minimap = self.user_settings.get('show_in_minimap') or self.settings.get('show_in_minimap')

        # Show information in status bar
        self.show_status = self.user_settings.get('show_status') or self.settings.get('show_status')
        if self.show_status != 'all' and self.show_status != 'none':
            self.show_status = 'default'
<commit_msg>Search for git executable in the PATH for ST2<commit_after>import os
import subprocess
import re
import codecs
import shutil

import sublime

ST3 = int(sublime.version()) >= 3006

try:
    from . import git_helper
    from .view_collection import ViewCollection
except (ImportError, ValueError):
    import git_helper
    from view_collection import ViewCollection

class GitGutterHandler:
    git_binary_path_error_shown = False
    git_binary_path_fallback = None

    def __init__(self, view):
        self.load_settings()
        self.view = view
        self.git_temp_file = ViewCollection.git_tmp_file(self.view)
        self.buf_temp_file = ViewCollection.buf_tmp_file(self.view)
        self.git_tree = None
        self.git_dir = None
        self.git_path = None

    def _get_view_encoding(self):
        # get encoding and clean it for python ex: "Western (ISO 8859-1)"
        # NOTE(maelnor): are we need regex here?
        pattern = re.compile(r'.+\((.*)\)')
        encoding = self.view.encoding()
        if encoding == "Undefined":
            encoding = self.view.settings().get('default_encoding')
        if pattern.match(encoding):
            encoding = pattern.sub(r'\1', encoding)

        encoding = encoding.replace('with BOM', '')
        encoding = encoding.replace('Windows', 'cp')
        encoding = encoding.replace('-', '_')
        encoding = encoding.replace(' ', '')

        # work around with ConvertToUTF8 plugin
        origin_encoding = self.view.settings().get('origin_encoding')
        return origin_encoding or encoding

    def on_disk(self):
        # if the view is saved to disk
        on_disk = self.view.file_name() is not None
        if on_disk:
            self.git_tree = self.git_tree or git_helper.git_tree(self.view)
            self.git_dir = self.git_dir or git_helper.git_dir(self.git_tree)
            self.git_path = self.git_path or git_helper.git_file_path(
                self.view, self.git_tree
            )
        return on_disk

    def reset(self):
        if self.on_disk() and self.git_path and self.view.window():
            self.view.window().run_command('git_gutter')

    def get_git_path(self):
        return self.git_path

    def update_buf_file(self):
        chars = self.view.size()
        region = sublime.Region(0, chars)

        # Try conversion
        try:
            contents = self.view.substr(
                region).encode(self._get_view_encoding())
        except UnicodeError:
            # Fallback to utf8-encoding
            contents = self.view.substr(region).encode('utf-8')
        except LookupError:
            # May encounter an encoding we don't have a codec for
            contents = self.view.substr(region).encode('utf-8')

        contents = contents.replace(b'\r\n', b'\n')
        contents = contents.replace(b'\r', b'\n')

        with open(self.buf_temp_file, 'wb') as f:
            if self.view.encoding() == "UTF-8 with BOM":
                f.write(codecs.BOM_UTF8)

            f.write(contents)

    def update_git_file(self):
        # the git repo won't change that often
        # so we can easily wait 5 seconds
        # between updates for performance
        if ViewCollection.git_time(self.view) > 5:
            with open(self.git_temp_file, 'w'):
                pass

            args = [
                self.git_binary_path,
                '--git-dir=' + self.git_dir,
                '--work-tree=' + self.git_tree,
                'show',
                ViewCollection.get_compare(self.view) + ':' + self.git_path,
            ]
            try:
                contents = self.run_command(args)
                contents = contents.replace(b'\r\n', b'\n')
                contents = contents.replace(b'\r', b'\n')
                with open(self.git_temp_file, 'wb') as f:
                    f.write(contents)

                ViewCollection.update_git_time(self.view)
            except Exception:
                pass

    def total_lines(self):
        chars = self.view.size()
        region = sublime.Region(0, chars)
        lines = self.view.lines(region)
        return len(lines)

    # Parse unified diff with 0 lines of context.
    # Hunk range info format:
    #   @@ -3,2 +4,0 @@
    #     Hunk originally starting at line 3, and occupying 2 lines, now
    #     starts at line 4, and occupies 0 lines, i.e. it was deleted.
    #   @@ -9 +10,2 @@
    #     Hunk size can be omitted, and defaults to one line.
    # Dealing with ambiguous hunks:
    #   "A\nB\n" -> "C\n"
    #   Was 'A' modified, and 'B' deleted? Or 'B' modified, 'A' deleted?
    #   Or both deleted? To minimize confusion, let's simply mark the
    #   hunk as modified.
    def process_diff(self, diff_str):
        inserted = []
        modified = []
        deleted = []
        hunk_re = '^@@ \-(\d+),?(\d*) \+(\d+),?(\d*) @@'
        hunks = re.finditer(hunk_re, diff_str, re.MULTILINE)
        for hunk in hunks:
            start = int(hunk.group(3))
            old_size = int(hunk.group(2) or 1)
            new_size = int(hunk.group(4) or 1)
            if not old_size:
                inserted += range(start, start + new_size)
            elif not new_size:
                deleted += [start + 1]
            else:
                modified += range(start, start + new_size)
        return (inserted, modified, deleted)

    def diff(self):
        if self.on_disk() and self.git_path:
            self.update_git_file()
            self.update_buf_file()
            args = [
                self.git_binary_path, 'diff', '-U0', '--no-color', '--no-index',
                self.ignore_whitespace,
                self.patience_switch,
                self.git_temp_file,
                self.buf_temp_file,
            ]
            args = list(filter(None, args))  # Remove empty args
            results = self.run_command(args)
            encoding = self._get_view_encoding()
            try:
                decoded_results = results.decode(encoding.replace(' ', ''))
            except UnicodeError:
                try:
                    decoded_results = results.decode("utf-8")
                except UnicodeDecodeError:
                    decoded_results = ""
            except LookupError:
                try:
                    decoded_results = codecs.decode(results)
                except UnicodeDecodeError:
                    decoded_results = ""
            return self.process_diff(decoded_results)
        else:
            return ([], [], [])

    def untracked(self):
        return self.handle_files([])

    def ignored(self):
        return self.handle_files(['-i'])

    def handle_files(self, additionnal_args):
        if self.on_disk() and self.git_path:
            args = [
                self.git_binary_path,
                '--git-dir=' + self.git_dir,
                '--work-tree=' + self.git_tree,
                'ls-files', '--other', '--exclude-standard',
            ] + additionnal_args + [
                os.path.join(self.git_tree, self.git_path),
            ]
            args = list(filter(None, args))  # Remove empty args
            results = self.run_command(args)
            encoding = self._get_view_encoding()
            try:
                decoded_results = results.decode(encoding.replace(' ', ''))
            except UnicodeError:
                decoded_results = results.decode("utf-8")
            return (decoded_results != "")
        else:
            return False

    def git_commits(self):
        args = [
            self.git_binary_path,
            '--git-dir=' + self.git_dir,
            '--work-tree=' + self.git_tree,
            'log', '--all',
            '--pretty=%s\a%h %an <%aE>\a%ad (%ar)',
            '--date=local', '--max-count=9000'
        ]
        results = self.run_command(args)
        return results

    def git_branches(self):
        args = [
            self.git_binary_path,
            '--git-dir=' + self.git_dir,
            '--work-tree=' + self.git_tree,
            'for-each-ref',
            '--sort=-committerdate',
            '--format=%(subject)\a%(refname)\a%(objectname)',
            'refs/heads/'
        ]
        results = self.run_command(args)
        return results

    def git_tags(self):
        args = [
            self.git_binary_path,
            '--git-dir=' + self.git_dir,
            '--work-tree=' + self.git_tree,
            'show-ref',
            '--tags',
            '--abbrev=7'
        ]
        results = self.run_command(args)
        return results

    def git_current_branch(self):
        args = [
            self.git_binary_path,
            '--git-dir=' + self.git_dir,
            '--work-tree=' + self.git_tree,
            'rev-parse',
            '--abbrev-ref',
            'HEAD'
        ]
        result = self.run_command(args)
        return result

    def run_command(self, args):
        startupinfo = None
        if os.name == 'nt':
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        proc = subprocess.Popen(args, stdout=subprocess.PIPE,
                                startupinfo=startupinfo, stderr=subprocess.PIPE)
        return proc.stdout.read()

    def load_settings(self):
        self.settings = sublime.load_settings('GitGutter.sublime-settings')
        self.user_settings = sublime.load_settings(
            'Preferences.sublime-settings')

        # Git Binary Setting
        git_binary_setting = self.user_settings.get("git_binary") or \
                             self.settings.get("git_binary")
        if isinstance(git_binary_setting, dict):
            self.git_binary_path = git_binary_setting.get(sublime.platform())
            if not self.git_binary_path:
                self.git_binary_path = git_binary_setting.get('default')
        else:
            self.git_binary_path = git_binary_setting

        if self.git_binary_path:
            self.git_binary_path = os.path.expandvars(self.git_binary_path)
        elif self.git_binary_path_fallback:
            self.git_binary_path = self.git_binary_path_fallback
        elif ST3:
            self.git_binary_path = shutil.which("git")
            GitGutterHandler.git_binary_path_fallback = self.git_binary_path
        else:
            git_exe = "git.exe" if sublime.platform() == "windows" else "git"
            for folder in os.environ["PATH"].split(os.pathsep):
                path = os.path.join(folder.strip('"'), git_exe)
                if os.path.isfile(path) and os.access(path, os.X_OK):
                    self.git_binary_path = path
                    GitGutterHandler.git_binary_path_fallback = path
                    break

        if not self.git_binary_path:
            if not GitGutterHandler.git_binary_path_error_shown:
                GitGutterHandler.git_binary_path_error_shown = True
                msg = ("Your Git binary cannot be found.  If it is installed, add it "
                       "to your PATH environment variable, or add a `git_binary` setting "
                       "in the `User/GitGutter.sublime-settings` file.")
                sublime.error_message(msg)
                raise ValueError("Git binary not found.")

        # Ignore White Space Setting
        self.ignore_whitespace = self.settings.get('ignore_whitespace')
        if self.ignore_whitespace == 'all':
            self.ignore_whitespace = '-w'
        elif self.ignore_whitespace == 'eol':
            self.ignore_whitespace = '--ignore-space-at-eol'
        else:
            self.ignore_whitespace = ''

        # Patience Setting
        self.patience_switch = ''
        patience = self.settings.get('patience')
        if patience:
            self.patience_switch = '--patience'

        # Untracked files
        self.show_untracked = self.settings.get(
            'show_markers_on_untracked_file')

        # Show in minimap
        self.show_in_minimap = self.user_settings.get('show_in_minimap') or self.settings.get('show_in_minimap')

        # Show information in status bar
        self.show_status = self.user_settings.get('show_status') or self.settings.get('show_status')
        if self.show_status != 'all' and self.show_status != 'none':
            self.show_status = 'default'
<|endoftext|>


<commit_before>
glashammer_version = '0.1.1'
<commit_msg>[project @ aafshar@gmail.com-20080804001131-11j2i6vajuuhf8vo] bumped version<commit_after>
glashammer_version = '0.1.2'
<|endoftext|>


<commit_before>from __future__ import absolute_import, division, print_function, unicode_literals
from xgboost.sklearn import XGBClassifier, XGBRegressor
import xgboost as xgb
import pandas as pd
import operator
import seaborn as sns
import numpy as np
from sklearn import grid_search, metrics
import unicodecsv as csv

class Xgb:
    def __init__(self, df, target_column='', id_column='', target_type='binary', categorical_columns=[], num_training_rounds=500, verbose=1, early_stopping_rounds=None):
        """
        input params:
        - df (DataFrame): dataframe of training data
        - target_column (string): name of target column
        - id_column (string): name of id column
        - target_type (string): 'linear' or 'binary'
        - categorical_columns (list): list of column names of categorical data. Will perform one-hot encoding
        - verbose (bool): verbosity of printouts
        """
        if type(df) == pd.core.frame.DataFrame:
            self.df = df
            self.early_stopping_rounds = early_stopping_rounds
            if target_column:
                self.target_column = target_column
                self.id_column = id_column
                self.target_type = target_type
                self.categorical_columns = categorical_columns
                self.verbose = verbose
                self.num_training_rounds = num_training_rounds
                # init the classifier
                if self.target_type == 'binary':
                    self.scoring = 'auc'
                    self.clf = XGBClassifier(
                        learning_rate =0.1,
                        n_estimators = num_training_rounds,
                        subsample = 0.8,
                        colsample_bytree = 0.8,
                        objective = 'binary:logistic',
                        scale_pos_weight = 1,
                        seed = 123)
                elif self.target_type == 'linear':
                    self.scoring = 'rmse'
                    self.clf = XGBRegressor(
                            n_estimators = num_training_rounds,
                            objective = 'reg:linear'
                            )
            else:
                print('please provide target column name')
        else:
            print('please provide pandas dataframe')

    def train(self):
        print('#### preprocessing ####')
        self.df = self.preprocess(self.df)

        print('#### training ####')
        self.predictors = [x for x in self.df.columns if x not in [self.target_column, self.id_column]]
        xgb_param = self.clf.get_xgb_params()

        xgtrain  = xgb.DMatrix(self.df[self.predictors], label=self.df[self.target_column], missing=np.nan)
        try:
            cvresult = xgb.cv(xgb_param, xgtrain, num_boost_round=self.clf.get_params()['n_estimators'], nfold=5,
                metrics=[self.scoring], early_stopping_rounds=self.early_stopping_rounds, show_progress=self.verbose)
        except:
            try:
                cvresult = xgb.cv(xgb_param, xgtrain, num_boost_round=self.clf.get_params()['n_estimators'], nfold=5,
                    metrics=[self.scoring], early_stopping_rounds=self.early_stopping_rounds, verbose_eval=self.verbose)
            except:
                cvresult = xgb.cv(xgb_param, xgtrain, num_boost_round=self.clf.get_params()['n_estimators'], nfold=5,
                    metrics=[self.scoring], early_stopping_rounds=self.early_stopping_rounds)
        self.clf.set_params(n_estimators=cvresult.shape[0])
        self.clf.fit(self.df[self.predictors], self.df[self.target_column],eval_metric=self.scoring)

        #Predict training set:
        train_df_predictions = self.clf.predict(self.df[self.predictors])

        if self.target_type == 'binary':
            train_df_predprob = self.clf.predict_proba(self.df[self.predictors])[:,1]
            print("Accuracy : %.4g" % metrics.accuracy_score(self.df[self.target_column].values, train_df_predictions))
            print("AUC Score (Train): %f" % metrics.roc_auc_score(self.df[self.target_column], train_df_predprob))
        elif self.target_type == 'linear':
            print("Mean squared error: %f" % metrics.mean_squared_error(self.df[self.target_column].values, train_df_predictions))
            print("Root mean squared error: %f" % np.sqrt(metrics.mean_squared_error(self.df[self.target_column].values, train_df_predictions)))

    def predict(self, test_df):
        print('### predicting ###')
        print('## preprocessing test set')
        self.test_df = self.preprocess(test_df, train=False)
        if self.target_type == 'binary':
            self.output = self.clf.predict_proba(self.test_df[self.predictors])[:,1]
        elif self.target_type == 'linear':
            self.output = self.clf.predict(self.test_df[self.predictors])
        return self.output

    def feature_importance(self, num_print=10, display=True):
        feature_importance = sorted(list(self.clf.booster().get_fscore().items()), key = operator.itemgetter(1), reverse=True)

        impt = pd.DataFrame(feature_importance)
        impt.columns = ['feature', 'importance']
        print(impt[:num_print])
        if display:
            impt[:num_print].plot("feature", "importance", kind="barh", color=sns.color_palette("deep", 3))

    def preprocess(self, df, train=True):
        # one hot encoding of categorical variables
        print('## one hot encoding of categorical variables')
        for col in self.categorical_columns:
            if self.verbose:
                print('one hot encoding: ', col)
            df = pd.concat([df, pd.get_dummies(df[col]).rename(columns=lambda x: col+'_'+str(x))], axis=1)
            df = df.drop([col], axis=1)

        # if training, determine columns to be removed
        if train:
            # drop columns that are too sparse to be informative
            self.cols_to_remove = []
            print('## dropping columns below sparsity threshold')
            for col in df.columns:
                nan_cnt = 0
                for x in df[col]:
                    try:
                        if np.isnan(x):
                            nan_cnt += 1
                    except:
                        pass
                if nan_cnt/float(len(df[col])) > 0.6: # arbitrary cutoff, if more than 60% missing then drop
                    if self.verbose:
                        print('will drop', col)
                    self.cols_to_remove.append(col)

            # drop columns that have no standard deviation (not informative)
            print('## dropping columns with no variation')
            for col in df.columns:
                if df[col].dtype == 'int64' or df[col].dtype == 'float64':
                    if df[col].std() == 0:
                        print('will drop', col)
                        self.cols_to_remove.append(col)
        if self.verbose and self.cols_to_remove:
            print('dropping the following columns:', self.cols_to_remove)
            df = df.drop(self.cols_to_remove, axis=1)

        if self.verbose:
            print('## DataFrame shape is now:', df.shape)

        # convert to numerical where possible
        print('## converting numerical data to numeric dtype')
        df = df.convert_objects(convert_numeric=True)

        # drop all those that are object type
        print('## dropping non-numerical columns')
        for col in df.columns:
            if df[col].dtype == 'int64' or df[col].dtype == 'float64' or df[col].dtype == 'bool':
                pass
            else:
                if self.verbose:
                    print('dropping because not int, float, or bool:', col)
                df = df.drop([col], axis=1)
        return df

    def write_csv(self, filename):
        """
        write results to csv
        """
        with open(filename, 'wb') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([self.id_column, self.target_column])
            for idx, value in enumerate(self.output):
                test_id = self.test_df['ID'][idx]
                test_output = self.output[idx]
                writer.writerow([test_id, test_output])

<commit_msg>add options to drop columns or convert to int or float. conversion to numeric no longer automatic<commit_after>from __future__ import absolute_import, division, print_function, unicode_literals
from xgboost.sklearn import XGBClassifier, XGBRegressor
import xgboost as xgb
import pandas as pd
import operator
import seaborn as sns
import numpy as np
from sklearn import grid_search, metrics
import unicodecsv as csv

class Xgb:
    def __init__(self, df, target_column='', id_column='', target_type='binary', categorical_columns=[], drop_columns=[], int_columns=[], float_columns=[], num_training_rounds=500, verbose=1, early_stopping_rounds=None):
        """
        input params:
        - df (DataFrame): dataframe of training data
        - target_column (string): name of target column
        - id_column (string): name of id column
        - target_type (string): 'linear' or 'binary'
        - categorical_columns (list): list of column names of categorical data. Will perform one-hot encoding
        - drop_columns (list): list of columns to drop
        - int_columns (list): list of columns to convert to int
        - float_columns (list): list of columns to convert to float
        - verbose (bool): verbosity of printouts
        """
        if type(df) == pd.core.frame.DataFrame:
            self.df = df
            self.early_stopping_rounds = early_stopping_rounds
            if target_column:
                self.target_column = target_column
                self.id_column = id_column
                self.target_type = target_type
                self.categorical_columns = categorical_columns
                self.verbose = verbose
                self.num_training_rounds = num_training_rounds
                # init the classifier
                if self.target_type == 'binary':
                    self.scoring = 'auc'
                    self.clf = XGBClassifier(
                        learning_rate =0.1,
                        n_estimators = num_training_rounds,
                        subsample = 0.8,
                        colsample_bytree = 0.8,
                        objective = 'binary:logistic',
                        scale_pos_weight = 1,
                        seed = 123)
                elif self.target_type == 'linear':
                    self.scoring = 'rmse'
                    self.clf = XGBRegressor(
                            n_estimators = num_training_rounds,
                            objective = 'reg:linear'
                            )
            else:
                print('please provide target column name')
        else:
            print('please provide pandas dataframe')

    def train(self):
        print('#### preprocessing ####')
        self.df = self.preprocess(self.df)

        print('#### training ####')
        self.predictors = [x for x in self.df.columns if x not in [self.target_column, self.id_column]]
        xgb_param = self.clf.get_xgb_params()

        xgtrain  = xgb.DMatrix(self.df[self.predictors], label=self.df[self.target_column], missing=np.nan)
        try:
            cvresult = xgb.cv(xgb_param, xgtrain, num_boost_round=self.clf.get_params()['n_estimators'], nfold=5,
                metrics=[self.scoring], early_stopping_rounds=self.early_stopping_rounds, show_progress=self.verbose)
        except:
            try:
                cvresult = xgb.cv(xgb_param, xgtrain, num_boost_round=self.clf.get_params()['n_estimators'], nfold=5,
                    metrics=[self.scoring], early_stopping_rounds=self.early_stopping_rounds, verbose_eval=self.verbose)
            except:
                cvresult = xgb.cv(xgb_param, xgtrain, num_boost_round=self.clf.get_params()['n_estimators'], nfold=5,
                    metrics=[self.scoring], early_stopping_rounds=self.early_stopping_rounds)
        self.clf.set_params(n_estimators=cvresult.shape[0])
        self.clf.fit(self.df[self.predictors], self.df[self.target_column],eval_metric=self.scoring)

        #Predict training set:
        train_df_predictions = self.clf.predict(self.df[self.predictors])

        if self.target_type == 'binary':
            train_df_predprob = self.clf.predict_proba(self.df[self.predictors])[:,1]
            print("Accuracy : %.4g" % metrics.accuracy_score(self.df[self.target_column].values, train_df_predictions))
            print("AUC Score (Train): %f" % metrics.roc_auc_score(self.df[self.target_column], train_df_predprob))
        elif self.target_type == 'linear':
            print("Mean squared error: %f" % metrics.mean_squared_error(self.df[self.target_column].values, train_df_predictions))
            print("Root mean squared error: %f" % np.sqrt(metrics.mean_squared_error(self.df[self.target_column].values, train_df_predictions)))

    def predict(self, test_df):
        print('### predicting ###')
        print('## preprocessing test set')
        self.test_df = self.preprocess(test_df, train=False)
        if self.target_type == 'binary':
            self.output = self.clf.predict_proba(self.test_df[self.predictors])[:,1]
        elif self.target_type == 'linear':
            self.output = self.clf.predict(self.test_df[self.predictors])
        return self.output

    def feature_importance(self, num_print=10, display=True):
        feature_importance = sorted(list(self.clf.booster().get_fscore().items()), key = operator.itemgetter(1), reverse=True)

        impt = pd.DataFrame(feature_importance)
        impt.columns = ['feature', 'importance']
        print(impt[:num_print])
        if display:
            impt[:num_print].plot("feature", "importance", kind="barh", color=sns.color_palette("deep", 3))

    def preprocess(self, df, train=True):
        # one hot encoding of categorical variables
        print('## one hot encoding of categorical variables')
        for col in self.categorical_columns:
            if self.verbose:
                print('one hot encoding: ', col)
            df = pd.concat([df, pd.get_dummies(df[col]).rename(columns=lambda x: col+'_'+str(x))], axis=1)
            df = df.drop([col], axis=1)

        # if training, determine columns to be removed
        if train:
            # drop columns that are too sparse to be informative
            self.cols_to_remove = []
            print('## dropping columns below sparsity threshold')
            for col in df.columns:
                nan_cnt = 0
                for x in df[col]:
                    try:
                        if np.isnan(x):
                            nan_cnt += 1
                    except:
                        pass
                if nan_cnt/float(len(df[col])) > 0.6: # arbitrary cutoff, if more than 60% missing then drop
                    if self.verbose:
                        print('will drop', col)
                    self.cols_to_remove.append(col)

            # drop columns that have no standard deviation (not informative)
            print('## dropping columns with no variation')
            for col in df.columns:
                if df[col].dtype == 'int64' or df[col].dtype == 'float64':
                    if df[col].std() == 0:
                        print('will drop', col)
                        self.cols_to_remove.append(col)
        if self.verbose and self.cols_to_remove:
            print('dropping the following columns:', self.cols_to_remove)
            df = df.drop(self.cols_to_remove, axis=1)

        if self.verbose:
            print('## DataFrame shape is now:', df.shape)

        # convert to numerical where possible
        #print('## converting numerical data to numeric dtype')
        #df = df.convert_objects(convert_numeric=True)

        # convert columns specified to be int and float
        df[int_columns] = df[int_columns].astype(int)
        df[float_columns] = df[float_columns].astype(float)

        # drop all those that are object type
        print('## dropping non-numerical columns')
        for col in df.columns:
            if df[col].dtype == 'int64' or df[col].dtype == 'float64' or df[col].dtype == 'bool':
                pass
            else:
                if self.verbose:
                    print('dropping because not int, float, or bool:', col)
                df = df.drop([col], axis=1)
        return df

    def write_csv(self, filename):
        """
        write results to csv
        """
        with open(filename, 'wb') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([self.id_column, self.target_column])
            for idx, value in enumerate(self.output):
                test_id = self.test_df['ID'][idx]
                test_output = self.output[idx]
                writer.writerow([test_id, test_output])

<|endoftext|>


<commit_before>"""
Xarray API for xhistogram.
"""

import xarray as xr
from collections import OrderedDict
from .core import histogram as _histogram

# range is a keyword so save the builtin so they can use it.
_range = range

def histogram(
    *args,
    bins=None,
    range=None,
    dim=None,
    weights=None,
    density=False,
    block_size="auto",
    keep_coords=False,
    bin_dim_suffix="_bin"
):
    """Histogram applied along specified dimensions.

    Parameters
    ----------
    args : xarray.DataArray objects
        Input data. The number of input arguments determines the dimensonality of
        the histogram. For example, two arguments prodocue a 2D histogram. All
        args must be aligned and have the same dimensions.
    bins :  int, str or numpy array or a list of ints, strs and/or arrays, optional
        If a list, there should be one entry for each item in ``args``.
        The bin specifications are as follows:

          * If int; the number of bins for all arguments in ``args``.
          * If str; the method used to automatically calculate the optimal bin width
            for all arguments in ``args``, as defined by numpy `histogram_bin_edges`.
          * If numpy array; the bin edges for all arguments in ``args``.
          * If a list of ints, strs and/or arrays; the bin specification as
            above for every argument in ``args``.

        When bin edges are specified, all but the last (righthand-most) bin include
        the left edge and exclude the right edge. The last bin includes both edges.

        A TypeError will be raised if args contains dask arrays and bins are not
        specified explicitly as an array or list of arrays. This is because other
        bin specifications trigger computation.
    range : (float, float) or a list of (float, float), optional
        If a list, there should be one entry for each item in ``args``.
        The range specifications are as follows:

          * If (float, float); the lower and upper range(s) of the bins for all
            arguments in ``args``. Values outside the range are ignored. The first
            element of the range must be less than or equal to the second. `range`
            affects the automatic bin computation as well. In this case, while bin
            width is computed to be optimal based on the actual data within `range`,
            the bin count will fill the entire range including portions containing
            no data.
          * If a list of (float, float); the ranges as above for every argument in
            ``args``.
          * If not provided, range is simply ``(arg.min(), arg.max())`` for each
            arg.
    dim : tuple of strings, optional
        Dimensions over which which the histogram is computed. The default is to
        compute the histogram of the flattened array.
    weights : array_like, optional
        An array of weights, of the same shape as `a`.  Each value in
        `a` only contributes its associated weight towards the bin count
        (instead of 1). If `density` is True, the weights are
        normalized, so that the integral of the density over the range
        remains 1. NaNs in the weights input will fill the entire bin with
        NaNs. If there are NaNs in the weights input call ``.fillna(0.)``
        before running ``histogram()``.
    density : bool, optional
        If ``False``, the result will contain the number of samples in
        each bin. If ``True``, the result is the value of the
        probability *density* function at the bin, normalized such that
        the *integral* over the range is 1. Note that the sum of the
        histogram values will not be equal to 1 unless bins of unity
        width are chosen; it is not a probability *mass* function.
    block_size : int or 'auto', optional
        A parameter which governs the algorithm used to compute the histogram.
        Using a nonzero value splits the histogram calculation over the
        non-histogram axes into blocks of size ``block_size``, iterating over
        them with a loop (numpy inputs) or in parallel (dask inputs). If
        ``'auto'``, blocks will be determined either by the underlying dask
        chunks (dask inputs) or an experimental built-in heuristic (numpy inputs).
    keep_coords : bool, optional
        If ``True``, keep all coordinates. Default: ``False``
    bin_dim_suffix : str, optional
        Suffix to append to input arg names to define names of output bin
        dimensions

    Returns
    -------
    hist : xarray.DataArray
        The values of the histogram. For each bin, the midpoint of the bin edges
        is given along the bin coordinates.

    """

    N_args = len(args)

    # TODO: allow list of weights as well
    N_weights = 1 if weights is not None else 0

    for a in args:
        # TODO: make this a more robust check
        assert a.name is not None, "all arrays must have a name"

    # we drop coords to simplify alignment
    if not keep_coords:
        args = [da.reset_coords(drop=True) for da in args]
    if N_weights:
        args += [weights.reset_coords(drop=True)]
    # explicitly broadcast so we understand what is going into apply_ufunc
    # (apply_ufunc might be doing this by itself again)
    args = list(xr.align(*args, join="exact"))

    # what happens if we skip this?
    # args = list(xr.broadcast(*args))
    a0 = args[0]
    a_coords = a0.coords

    # roll our own broadcasting
    # now manually expand the arrays
    all_dims = [d for a in args for d in a.dims]
    all_dims_ordered = list(OrderedDict.fromkeys(all_dims))
    args_expanded = []
    for a in args:
        expand_keys = [d for d in all_dims_ordered if d not in a.dims]
        a_expanded = a.expand_dims({k: 1 for k in expand_keys})
        args_expanded.append(a_expanded)

    # only transpose if necessary, to avoid creating unnecessary dask tasks
    args_transposed = []
    for a in args_expanded:
        if a.dims != all_dims_ordered:
            args_transposed.append(a.transpose(*all_dims_ordered))
        else:
            args.transposed.append(a)
    args_data = [a.data for a in args_transposed]

    if N_weights:
        weights_data = args_data.pop()
    else:
        weights_data = None

    if dim is not None:
        dims_to_keep = [d for d in all_dims_ordered if d not in dim]
        axis = [args_transposed[0].get_axis_num(d) for d in dim]
    else:
        dims_to_keep = []
        axis = None

    h_data, bins = _histogram(
        *args_data,
        weights=weights_data,
        bins=bins,
        range=range,
        axis=axis,
        density=density,
        block_size=block_size
    )

    # create output dims
    new_dims = [a.name + bin_dim_suffix for a in args[:N_args]]
    output_dims = dims_to_keep + new_dims

    # create new coords
    bin_centers = [0.5 * (bin[:-1] + bin[1:]) for bin in bins]
    new_coords = {
        name: ((name,), bin_center, a.attrs)
        for name, bin_center, a in zip(new_dims, bin_centers, args)
    }

    # old coords associated with dims
    old_dim_coords = {name: a0[name] for name in dims_to_keep if name in a_coords}

    all_coords = {}
    all_coords.update(old_dim_coords)
    all_coords.update(new_coords)
    # add compatible coords
    if keep_coords:
        for c in a_coords:
            if c not in all_coords and set(a0[c].dims).issubset(output_dims):
                all_coords[c] = a0[c]

    output_name = "_".join(["histogram"] + [a.name for a in args[:N_args]])

    da_out = xr.DataArray(h_data, dims=output_dims, coords=all_coords, name=output_name)

    if density:
        # correct for overcounting the bins which weren't histogrammed along
        n_bins_bystander_dims = da_out.isel(**{bd: 0 for bd in new_dims}).size
        da_out = da_out * n_bins_bystander_dims

    return da_out

    # we need weights to be passed through apply_func's alignment algorithm,
    # so we include it as an arg, so we create a wrapper function to do so
    # this feels like a hack
    # def _histogram_wrapped(*args, **kwargs):
    #     alist = list(args)
    #     weights = [alist.pop() for n in _range(N_weights)]
    #     if N_weights == 0:
    #         weights = None
    #     elif N_weights == 1:
    #         weights = weights[0] # squeeze
    #     return _histogram(*alist, weights=weights, **kwargs)
<commit_msg>Created checking function for xarray.DataArray and throw an error if not<commit_after>"""
Xarray API for xhistogram.
"""

import xarray as xr
from collections import OrderedDict
from .core import histogram as _histogram

# range is a keyword so save the builtin so they can use it.
_range = range

def histogram(
    *args,
    bins=None,
    range=None,
    dim=None,
    weights=None,
    density=False,
    block_size="auto",
    keep_coords=False,
    bin_dim_suffix="_bin"
):
    """Histogram applied along specified dimensions.

    Parameters
    ----------
    args : xarray.DataArray objects
        Input data. The number of input arguments determines the dimensonality of
        the histogram. For example, two arguments prodocue a 2D histogram. All
        args must be aligned and have the same dimensions.
    bins :  int, str or numpy array or a list of ints, strs and/or arrays, optional
        If a list, there should be one entry for each item in ``args``.
        The bin specifications are as follows:

          * If int; the number of bins for all arguments in ``args``.
          * If str; the method used to automatically calculate the optimal bin width
            for all arguments in ``args``, as defined by numpy `histogram_bin_edges`.
          * If numpy array; the bin edges for all arguments in ``args``.
          * If a list of ints, strs and/or arrays; the bin specification as
            above for every argument in ``args``.

        When bin edges are specified, all but the last (righthand-most) bin include
        the left edge and exclude the right edge. The last bin includes both edges.

        A TypeError will be raised if args contains dask arrays and bins are not
        specified explicitly as an array or list of arrays. This is because other
        bin specifications trigger computation.
    range : (float, float) or a list of (float, float), optional
        If a list, there should be one entry for each item in ``args``.
        The range specifications are as follows:

          * If (float, float); the lower and upper range(s) of the bins for all
            arguments in ``args``. Values outside the range are ignored. The first
            element of the range must be less than or equal to the second. `range`
            affects the automatic bin computation as well. In this case, while bin
            width is computed to be optimal based on the actual data within `range`,
            the bin count will fill the entire range including portions containing
            no data.
          * If a list of (float, float); the ranges as above for every argument in
            ``args``.
          * If not provided, range is simply ``(arg.min(), arg.max())`` for each
            arg.
    dim : tuple of strings, optional
        Dimensions over which which the histogram is computed. The default is to
        compute the histogram of the flattened array.
    weights : array_like, optional
        An array of weights, of the same shape as `a`.  Each value in
        `a` only contributes its associated weight towards the bin count
        (instead of 1). If `density` is True, the weights are
        normalized, so that the integral of the density over the range
        remains 1. NaNs in the weights input will fill the entire bin with
        NaNs. If there are NaNs in the weights input call ``.fillna(0.)``
        before running ``histogram()``.
    density : bool, optional
        If ``False``, the result will contain the number of samples in
        each bin. If ``True``, the result is the value of the
        probability *density* function at the bin, normalized such that
        the *integral* over the range is 1. Note that the sum of the
        histogram values will not be equal to 1 unless bins of unity
        width are chosen; it is not a probability *mass* function.
    block_size : int or 'auto', optional
        A parameter which governs the algorithm used to compute the histogram.
        Using a nonzero value splits the histogram calculation over the
        non-histogram axes into blocks of size ``block_size``, iterating over
        them with a loop (numpy inputs) or in parallel (dask inputs). If
        ``'auto'``, blocks will be determined either by the underlying dask
        chunks (dask inputs) or an experimental built-in heuristic (numpy inputs).
    keep_coords : bool, optional
        If ``True``, keep all coordinates. Default: ``False``
    bin_dim_suffix : str, optional
        Suffix to append to input arg names to define names of output bin
        dimensions

    Returns
    -------
    hist : xarray.DataArray
        The values of the histogram. For each bin, the midpoint of the bin edges
        is given along the bin coordinates.

    """

    N_args = len(args)

    # TODO: allow list of weights as well
    N_weights = 1 if weights is not None else 0

    for a in args:
        if not isinstance(a, xr.DataArray):
            raise TypeError(f"xhist() accepts an xr.DataArray object but a " \
                            f"{type(a).__name__} was given")

    # we drop coords to simplify alignment
    if not keep_coords:
        args = [da.reset_coords(drop=True) for da in args]
    if N_weights:
        args += [weights.reset_coords(drop=True)]
    # explicitly broadcast so we understand what is going into apply_ufunc
    # (apply_ufunc might be doing this by itself again)
    args = list(xr.align(*args, join="exact"))

    # what happens if we skip this?
    # args = list(xr.broadcast(*args))
    a0 = args[0]
    a_coords = a0.coords

    # roll our own broadcasting
    # now manually expand the arrays
    all_dims = [d for a in args for d in a.dims]
    all_dims_ordered = list(OrderedDict.fromkeys(all_dims))
    args_expanded = []
    for a in args:
        expand_keys = [d for d in all_dims_ordered if d not in a.dims]
        a_expanded = a.expand_dims({k: 1 for k in expand_keys})
        args_expanded.append(a_expanded)

    # only transpose if necessary, to avoid creating unnecessary dask tasks
    args_transposed = []
    for a in args_expanded:
        if a.dims != all_dims_ordered:
            args_transposed.append(a.transpose(*all_dims_ordered))
        else:
            args.transposed.append(a)
    args_data = [a.data for a in args_transposed]

    if N_weights:
        weights_data = args_data.pop()
    else:
        weights_data = None

    if dim is not None:
        dims_to_keep = [d for d in all_dims_ordered if d not in dim]
        axis = [args_transposed[0].get_axis_num(d) for d in dim]
    else:
        dims_to_keep = []
        axis = None

    h_data, bins = _histogram(
        *args_data,
        weights=weights_data,
        bins=bins,
        range=range,
        axis=axis,
        density=density,
        block_size=block_size
    )

    # create output dims
    new_dims = [a.name + bin_dim_suffix for a in args[:N_args]]
    output_dims = dims_to_keep + new_dims

    # create new coords
    bin_centers = [0.5 * (bin[:-1] + bin[1:]) for bin in bins]
    new_coords = {
        name: ((name,), bin_center, a.attrs)
        for name, bin_center, a in zip(new_dims, bin_centers, args)
    }

    # old coords associated with dims
    old_dim_coords = {name: a0[name] for name in dims_to_keep if name in a_coords}

    all_coords = {}
    all_coords.update(old_dim_coords)
    all_coords.update(new_coords)
    # add compatible coords
    if keep_coords:
        for c in a_coords:
            if c not in all_coords and set(a0[c].dims).issubset(output_dims):
                all_coords[c] = a0[c]

    output_name = "_".join(["histogram"] + [a.name for a in args[:N_args]])

    da_out = xr.DataArray(h_data, dims=output_dims, coords=all_coords, name=output_name)

    if density:
        # correct for overcounting the bins which weren't histogrammed along
        n_bins_bystander_dims = da_out.isel(**{bd: 0 for bd in new_dims}).size
        da_out = da_out * n_bins_bystander_dims

    return da_out

    # we need weights to be passed through apply_func's alignment algorithm,
    # so we include it as an arg, so we create a wrapper function to do so
    # this feels like a hack
    # def _histogram_wrapped(*args, **kwargs):
    #     alist = list(args)
    #     weights = [alist.pop() for n in _range(N_weights)]
    #     if N_weights == 0:
    #         weights = None
    #     elif N_weights == 1:
    #         weights = weights[0] # squeeze
    #     return _histogram(*alist, weights=weights, **kwargs)
<|endoftext|>


<commit_before>#!/usr/bin/python
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import sys, getopt, json, os, base64
from fcntl import flock, LOCK_EX, LOCK_UN

def main(argv):
    fpath =  ''
    b64data = ''

    try:
        opts, args = getopt.getopt(argv,"f:d:")
    except getopt.GetoptError:
        print 'params: -f <filename> -d <b64jsondata>'
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-f':
            fpath = arg
        elif opt == '-d':
            b64data = arg

    json_data = ''
    if fpath != '':
        fh = open(fpath, 'r')
        json_data = json.loads(fh.read())
    elif b64data != '':
        json_data = json.loads(base64.b64decode(b64data))
    else:
        print '-f <filename> or -d <b64jsondata> required'
        sys.exit(2)

    for ip in json_data:
        for item in json_data[ip]:
            folder = item[0]
            file   = item[1]
            data   = item[2]

            # process only valid data
            if folder != "userdata" and folder != "metadata":
                continue

            if file == "":
                continue

            htaccess(ip, folder, file)

            if data == "":
                deletefile(ip, folder, file)
            else:
                createfile(ip, folder, file, data)

    if fpath != '':
        fh.close()
        os.remove(fpath)

def deletefile(ip, folder, file):
    datafile = "/var/www/html/" + folder + "/" + ip + "/" + file

    if os.path.exists(datafile):
        os.remove(datafile)

def createfile(ip, folder, file, data):
    dest = "/var/www/html/" + folder + "/" + ip + "/" + file
    metamanifestdir = "/var/www/html/" + folder + "/" + ip
    metamanifest =  metamanifestdir + "/meta-data"

    # base64 decode userdata
    if folder == "userdata" or folder == "user-data":
        if data is not None:
            data = base64.b64decode(data)

    fh = open(dest, "w")
    exflock(fh)
    if data is not None:
        fh.write(data)
    else:
        fh.write("")
    unflock(fh)
    fh.close()
    os.chmod(dest, 0644)

    if folder == "metadata" or folder == "meta-data":
        try:
            os.makedirs(metamanifestdir, 0755)
        except OSError as e:
            # error 17 is already exists, we do it this way for concurrency
            if e.errno != 17:
                print "failed to make directories " + metamanifestdir + " due to :" +e.strerror
                sys.exit(1)
        if os.path.exists(metamanifest):
            fh = open(metamanifest, "r+a")
            exflock(fh)
            if not file in fh.read():
                fh.write(file + '\n')
            unflock(fh)
            fh.close()
        else:
            fh = open(metamanifest, "w")
            exflock(fh)
            fh.write(file + '\n')
            unflock(fh)
            fh.close()

    if os.path.exists(metamanifest):
        os.chmod(metamanifest, 0644)

def htaccess(ip, folder, file):
    entry = "RewriteRule ^" + file + "$  ../" + folder + "/%{REMOTE_ADDR}/" + file + " [L,NC,QSA]"
    htaccessFolder = "/var/www/html/latest"
    htaccessFile = htaccessFolder + "/.htaccess"

    try:
        os.mkdir(htaccessFolder,0755)
    except OSError as e:
        # error 17 is already exists, we do it this way for concurrency
        if e.errno != 17:
            print "failed to make directories " + htaccessFolder + " due to :" +e.strerror
            sys.exit(1)

    if os.path.exists(htaccessFile):
        fh = open(htaccessFile, "r+a")
        exflock(fh)
        if not entry in fh.read():
            fh.write(entry + '\n')
        unflock(fh)
        fh.close()
    else:
        fh = open(htaccessFile, "w")
        exflock(fh)
        fh.write("Options +FollowSymLinks\nRewriteEngine On\n\n")
        fh.write(entry + '\n')
        unflock(fh)
        fh.close()

    entry="Options -Indexes\nOrder Deny,Allow\nDeny from all\nAllow from " + ip
    htaccessFolder = "/var/www/html/" + folder + "/" + ip
    htaccessFile = htaccessFolder+"/.htaccess"

    try:
        os.makedirs(htaccessFolder,0755)
    except OSError as e:
        # error 17 is already exists, we do it this way for sake of concurrency
        if e.errno != 17:
            print "failed to make directories " + htaccessFolder + " due to :" +e.strerror
            sys.exit(1)

    fh = open(htaccessFile, "w")
    exflock(fh)
    fh.write(entry + '\n')
    unflock(fh)
    fh.close()

    if folder == "metadata" or folder == "meta-data":
        entry = "RewriteRule ^meta-data/(.+)$  ../" + folder + "/%{REMOTE_ADDR}/$1 [L,NC,QSA]"
        htaccessFolder = "/var/www/html/latest"
        htaccessFile = htaccessFolder + "/.htaccess"

        fh = open(htaccessFile, "r+a")
        exflock(fh)
        if not entry in fh.read():
            fh.write(entry + '\n')

        entry = "RewriteRule ^meta-data/$  ../" + folder + "/%{REMOTE_ADDR}/meta-data [L,NC,QSA]"

        fh.seek(0)
        if not entry in fh.read():
            fh.write(entry + '\n')
        unflock(fh)
        fh.close()

def exflock(file):
    try:
        flock(file, LOCK_EX)
    except IOError as e:
        print "failed to lock file" + file.name + " due to : " + e.strerror
        sys.exit(1)
    return True

def unflock(file):
    try:
        flock(file, LOCK_UN)
    except IOError:
        print "failed to unlock file" + file.name + " due to : " + e.strerror
        sys.exit(1)
    return True

if __name__ == "__main__":
    main(sys.argv[1:])
<commit_msg>CLOUDSTACK-7405: Make the trailing slash conditional for meta-data requests<commit_after>#!/usr/bin/python
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import sys, getopt, json, os, base64
from fcntl import flock, LOCK_EX, LOCK_UN

def main(argv):
    fpath =  ''
    b64data = ''

    try:
        opts, args = getopt.getopt(argv,"f:d:")
    except getopt.GetoptError:
        print 'params: -f <filename> -d <b64jsondata>'
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-f':
            fpath = arg
        elif opt == '-d':
            b64data = arg

    json_data = ''
    if fpath != '':
        fh = open(fpath, 'r')
        json_data = json.loads(fh.read())
    elif b64data != '':
        json_data = json.loads(base64.b64decode(b64data))
    else:
        print '-f <filename> or -d <b64jsondata> required'
        sys.exit(2)

    for ip in json_data:
        for item in json_data[ip]:
            folder = item[0]
            file   = item[1]
            data   = item[2]

            # process only valid data
            if folder != "userdata" and folder != "metadata":
                continue

            if file == "":
                continue

            htaccess(ip, folder, file)

            if data == "":
                deletefile(ip, folder, file)
            else:
                createfile(ip, folder, file, data)

    if fpath != '':
        fh.close()
        os.remove(fpath)

def deletefile(ip, folder, file):
    datafile = "/var/www/html/" + folder + "/" + ip + "/" + file

    if os.path.exists(datafile):
        os.remove(datafile)

def createfile(ip, folder, file, data):
    dest = "/var/www/html/" + folder + "/" + ip + "/" + file
    metamanifestdir = "/var/www/html/" + folder + "/" + ip
    metamanifest =  metamanifestdir + "/meta-data"

    # base64 decode userdata
    if folder == "userdata" or folder == "user-data":
        if data is not None:
            data = base64.b64decode(data)

    fh = open(dest, "w")
    exflock(fh)
    if data is not None:
        fh.write(data)
    else:
        fh.write("")
    unflock(fh)
    fh.close()
    os.chmod(dest, 0644)

    if folder == "metadata" or folder == "meta-data":
        try:
            os.makedirs(metamanifestdir, 0755)
        except OSError as e:
            # error 17 is already exists, we do it this way for concurrency
            if e.errno != 17:
                print "failed to make directories " + metamanifestdir + " due to :" +e.strerror
                sys.exit(1)
        if os.path.exists(metamanifest):
            fh = open(metamanifest, "r+a")
            exflock(fh)
            if not file in fh.read():
                fh.write(file + '\n')
            unflock(fh)
            fh.close()
        else:
            fh = open(metamanifest, "w")
            exflock(fh)
            fh.write(file + '\n')
            unflock(fh)
            fh.close()

    if os.path.exists(metamanifest):
        os.chmod(metamanifest, 0644)

def htaccess(ip, folder, file):
    entry = "RewriteRule ^" + file + "$  ../" + folder + "/%{REMOTE_ADDR}/" + file + " [L,NC,QSA]"
    htaccessFolder = "/var/www/html/latest"
    htaccessFile = htaccessFolder + "/.htaccess"

    try:
        os.mkdir(htaccessFolder,0755)
    except OSError as e:
        # error 17 is already exists, we do it this way for concurrency
        if e.errno != 17:
            print "failed to make directories " + htaccessFolder + " due to :" +e.strerror
            sys.exit(1)

    if os.path.exists(htaccessFile):
        fh = open(htaccessFile, "r+a")
        exflock(fh)
        if not entry in fh.read():
            fh.write(entry + '\n')
        unflock(fh)
        fh.close()
    else:
        fh = open(htaccessFile, "w")
        exflock(fh)
        fh.write("Options +FollowSymLinks\nRewriteEngine On\n\n")
        fh.write(entry + '\n')
        unflock(fh)
        fh.close()

    entry="Options -Indexes\nOrder Deny,Allow\nDeny from all\nAllow from " + ip
    htaccessFolder = "/var/www/html/" + folder + "/" + ip
    htaccessFile = htaccessFolder+"/.htaccess"

    try:
        os.makedirs(htaccessFolder,0755)
    except OSError as e:
        # error 17 is already exists, we do it this way for sake of concurrency
        if e.errno != 17:
            print "failed to make directories " + htaccessFolder + " due to :" +e.strerror
            sys.exit(1)

    fh = open(htaccessFile, "w")
    exflock(fh)
    fh.write(entry + '\n')
    unflock(fh)
    fh.close()

    if folder == "metadata" or folder == "meta-data":
        entry = "RewriteRule ^meta-data/(.+)$  ../" + folder + "/%{REMOTE_ADDR}/$1 [L,NC,QSA]"
        htaccessFolder = "/var/www/html/latest"
        htaccessFile = htaccessFolder + "/.htaccess"

        fh = open(htaccessFile, "r+a")
        exflock(fh)
        if not entry in fh.read():
            fh.write(entry + '\n')

        entry = "RewriteRule ^meta-data/?$  ../" + folder + "/%{REMOTE_ADDR}/meta-data [L,NC,QSA]"

        fh.seek(0)
        if not entry in fh.read():
            fh.write(entry + '\n')
        unflock(fh)
        fh.close()

def exflock(file):
    try:
        flock(file, LOCK_EX)
    except IOError as e:
        print "failed to lock file" + file.name + " due to : " + e.strerror
        sys.exit(1)
    return True

def unflock(file):
    try:
        flock(file, LOCK_UN)
    except IOError:
        print "failed to unlock file" + file.name + " due to : " + e.strerror
        sys.exit(1)
    return True

if __name__ == "__main__":
    main(sys.argv[1:])
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing,
#  software distributed under the License is distributed on an
#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#  KIND, either express or implied.  See the License for the
#  specific language governing permissions and limitations
#  under the License.

import os
import re
import json
import copy
import time
from collections import defaultdict, OrderedDict

import requests
import _jsonnet  # pylint: disable=import-error

LINUX_PRICING_URLS = [
    # Deprecated instances (JSON format)
    "https://aws.amazon.com/ec2/pricing/json/linux-od.json",
    # Previous generation instances (JavaScript file)
    "https://a0.awsstatic.com/pricing/1/ec2/previous-generation/linux-od.min.js",
    # New generation instances (JavaScript file)
    # Using other endpoint atm
    # 'https://a0.awsstatic.com/pricing/1/ec2/linux-od.min.js'
]

EC2_REGIONS = [
    "us-east-1",
    "us-east-2",
    "us-west-1",
    "us-west-2",
    "us-gov-west-1",
    "eu-west-1",
    "eu-west-2",
    "eu-west-3",
    "eu-north-1",
    "eu-south-1",
    "eu-central-1",
    "ca-central-1",
    "ap-southeast-1",
    "ap-southeast-2",
    "ap-northeast-1",
    "ap-northeast-2",
    "ap-south-1",
    "sa-east-1",
    "cn-north-1",
    "ap-east-1",
    "af-south-1",
]

EC2_INSTANCE_TYPES = [
    "t1.micro",
    "m1.small",
    "m1.medium",
    "m1.large",
    "m1.xlarge",
    "m2.xlarge",
    "m2.2xlarge",
    "m2.4xlarge",
    "m3.medium",
    "m3.large",
    "m3.xlarge",
    "m3.2xlarge",
    "c1.medium",
    "c1.xlarge",
    "cc1.4xlarge",
    "cc2.8xlarge",
    "c3.large",
    "c3.xlarge",
    "c3.2xlarge",
    "c3.4xlarge",
    "c3.8xlarge",
    "d2.xlarge",
    "d2.2xlarge",
    "d2.4xlarge",
    "d2.8xlarge",
    "cg1.4xlarge",
    "g2.2xlarge",
    "g2.8xlarge",
    "cr1.8xlarge",
    "hs1.4xlarge",
    "hs1.8xlarge",
    "i2.xlarge",
    "i2.2xlarge",
    "i2.4xlarge",
    "i2.8xlarge",
    "i3.large",
    "i3.xlarge",
    "i3.2xlarge",
    "i3.4xlarge",
    "i3.8xlarge",
    "i3.16large",
    "r3.large",
    "r3.xlarge",
    "r3.2xlarge",
    "r3.4xlarge",
    "r3.8xlarge",
    "r4.large",
    "r4.xlarge",
    "r4.2xlarge",
    "r4.4xlarge",
    "r4.8xlarge",
    "r4.16xlarge",
    "t2.micro",
    "t2.small",
    "t2.medium",
    "t2.large",
    "x1.32xlarge",
]

# Maps EC2 region name to region name used in the pricing file
REGION_NAME_MAP = {
    "us-east": "ec2_us_east",
    "us-east-1": "ec2_us_east",
    "us-east-2": "ec2_us_east_ohio",
    "us-west": "ec2_us_west",
    "us-west-1": "ec2_us_west",
    "us-west-2": "ec2_us_west_oregon",
    "eu-west-1": "ec2_eu_west",
    "eu-west-2": "ec2_eu_west_london",
    "eu-west-3": "ec2_eu_west_3",
    "eu-ireland": "ec2_eu_west",
    "eu-south-1": "ec2_eu_south",
    "eu-central-1": "ec2_eu_central",
    "ca-central-1": "ec2_ca_central_1",
    "apac-sin": "ec2_ap_southeast",
    "ap-southeast-1": "ec2_ap_southeast",
    "apac-syd": "ec2_ap_southeast_2",
    "ap-southeast-2": "ec2_ap_southeast_2",
    "apac-tokyo": "ec2_ap_northeast",
    "ap-northeast-1": "ec2_ap_northeast",
    "ap-northeast-2": "ec2_ap_northeast",
    "ap-south-1": "ec2_ap_south_1",
    "sa-east-1": "ec2_sa_east",
    "us-gov-west-1": "ec2_us_govwest",
    "cn-north-1": "ec2_cn_north",
    "ap-east-1": "ec2_ap_east",
    "af-south-1": "ec2_af_south",
}

INSTANCE_SIZES = [
    "micro",
    "small",
    "medium",
    "large",
    "xlarge",
    "x-large",
    "extra-large",
]

RE_NUMERIC_OTHER = re.compile(r"(?:([0-9]+)|([-A-Z_a-z]+)|([^-0-9A-Z_a-z]+))")

BASE_PATH = os.path.dirname(os.path.abspath(__file__))
PRICING_FILE_PATH = os.path.join(BASE_PATH, "../libcloud/data/pricing.json")
PRICING_FILE_PATH = os.path.abspath(PRICING_FILE_PATH)

def scrape_ec2_pricing():
    result = defaultdict(OrderedDict)
    os_map = {"linux": "ec2_linux", "windows-std": "ec2_windows"}
    for item in os_map.values():
        result[item] = {}
    for url in LINUX_PRICING_URLS:
        response = requests.get(url)

        if re.match(r".*?\.json$", url):
            data = response.json()
            print("Sample response: %s..." % (str(data)[:100]))
        elif re.match(r".*?\.js$", url):
            data = response.content.decode("utf-8")
            print("Sample response: %s..." % (data[:100]))
            match = re.match(r"^.*callback\((.*?)\);?$", data, re.MULTILINE | re.DOTALL)
            data = match.group(1)
            # NOTE: We used to use demjson, but it's not working under Python 3 and new version of
            # setuptools anymore so we use jsonnet
            # demjson supports non-strict mode and can parse unquoted objects
            data = json.loads(_jsonnet.evaluate_snippet("snippet", data))
        regions = data["config"]["regions"]

        for region_data in regions:
            region_name = region_data["region"]
            instance_types = region_data["instanceTypes"]

            for instance_type in instance_types:
                sizes = instance_type["sizes"]
                for size in sizes:
                    if not result["ec2_linux"].get(size["size"], False):
                        result["ec2_linux"][size["size"]] = {}
                    price = size["valueColumns"][0]["prices"]["USD"]
                    if str(price).lower() == "n/a":
                        # Price not available
                        continue

                    result["ec2_linux"][size["size"]][region_name] = float(price)

    res = defaultdict(OrderedDict)
    url = "https://calculator.aws/pricing/1.0/" "ec2/region/{}/ondemand/{}/index.json"
    instances = set()
    for OS in ["linux", "windows-std"]:
        res[os_map[OS]] = {}
        for region in EC2_REGIONS:
            res[os_map[OS]][region] = {}
            full_url = url.format(region, OS)
            response = requests.get(full_url)
            if response.status_code != 200:
                print(
                    "Skipping URL %s which returned non 200-status code (%s)"
                    % (full_url, response.status_code)
                )
                continue
            data = response.json()

            for entry in data["prices"]:
                instance_type = entry["attributes"].get("aws:ec2:instanceType", "")
                instances.add(instance_type)
                price = entry["price"].get("USD", 0)
                res[os_map[OS]][region][instance_type] = price
    for item in os_map.values():
        for instance in instances:
            if not result[item].get(instance, False):
                result[item][instance] = {}
            for region in EC2_REGIONS:
                if res[item][region].get(instance, False):
                    result[item][instance][region] = float(res[item][region][instance])

    return result

def update_pricing_file(pricing_file_path, pricing_data):
    with open(pricing_file_path, "r") as fp:
        content = fp.read()

    data = json.loads(content)
    original_data = copy.deepcopy(data)

    data["compute"].update(pricing_data)

    if data == original_data:
        # Nothing has changed, bail out early and don't update "updated" attribute
        print("Nothing has changed, skipping update.")
        return

    data["updated"] = int(time.time())

    # Always sort the pricing info
    data = sort_nested_dict(data)

    content = json.dumps(data, indent=4)
    lines = content.splitlines()
    lines = [line.rstrip() for line in lines]
    content = "\n".join(lines)

    with open(pricing_file_path, "w") as fp:
        fp.write(content)

def sort_nested_dict(value):
    """
    Recursively sort a nested dict.
    """
    result = OrderedDict()

    for key, value in sorted(value.items(), key=sort_key_by_numeric_other):
        if isinstance(value, (dict, OrderedDict)):
            result[key] = sort_nested_dict(value)
        else:
            result[key] = value

    return result

def sort_key_by_numeric_other(key_value):
    """
    Split key into numeric, alpha and other part and sort accordingly.
    """
    result = []

    for (numeric, alpha, other) in RE_NUMERIC_OTHER.findall(key_value[0]):
        numeric = int(numeric) if numeric else -1
        alpha = INSTANCE_SIZES.index(alpha) if alpha in INSTANCE_SIZES else alpha
        alpha = str(alpha)
        item = tuple([numeric, alpha, other])
        result.append(item)

    return tuple(result)

def main():
    print("Scraping EC2 pricing data (this may take up to 2 minutes)....")

    pricing_data = scrape_ec2_pricing()
    update_pricing_file(pricing_file_path=PRICING_FILE_PATH, pricing_data=pricing_data)

    print("Pricing data updated")

if __name__ == "__main__":
    main()
<commit_msg>Use official endpoint for EC2 prices<commit_after>#!/usr/bin/env python
#
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing,
#  software distributed under the License is distributed on an
#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#  KIND, either express or implied.  See the License for the
#  specific language governing permissions and limitations
#  under the License.

import os
import re
import json
import copy
import time
from collections import defaultdict, OrderedDict

import requests
import ijson  # pylint: disable=import-error

# same URL as the one used by scrape-ec2-sizes.py, now it has official data on pricing
URL = "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonEC2/current/index.json"

RE_NUMERIC_OTHER = re.compile(r"(?:([0-9]+)|([-A-Z_a-z]+)|([^-0-9A-Z_a-z]+))")

BASE_PATH = os.path.dirname(os.path.abspath(__file__))
PRICING_FILE_PATH = os.path.join(BASE_PATH, "../libcloud/data/pricing.json")
PRICING_FILE_PATH = os.path.abspath(PRICING_FILE_PATH)

TEMPFILE = os.environ.get("TMP_JSON", "/tmp/ec.json")

INSTANCE_SIZES = [
    "micro",
    "small",
    "medium",
    "large",
    "xlarge",
    "x-large",
    "extra-large",
]

def download_json():
    response = requests.get(URL, stream=True)
    try:
        return open(TEMPFILE, "r")
    except IOError:
        with open(TEMPFILE, "wb") as fo:
            for chunk in response.iter_content(chunk_size=2**20):
                if chunk:
                    fo.write(chunk)
    return open(TEMPFILE, "r")

def get_json():
    try:
        return open(TEMPFILE, "r")
    except IOError:
        return download_json()

# Prices and sizes are in different dicts and categorized by sku
def get_all_prices():
    # return variable
    # prices = {sku : {price: int, unit: string}}
    prices = {}
    current_sku = ""
    current_rate_code = ""
    amazonEC2_offer_code = "JRTCKXETXF"
    json_file = get_json()
    parser = ijson.parse(json_file)
    # use parser because file is very large
    for prefix, event, value in parser:
        if "products" in prefix:
            continue
        if (prefix, event) == ("terms.OnDemand", "map_key"):
            current_sku = value
            prices[current_sku] = {}
        elif (prefix, event) == (
            f"terms.OnDemand.{current_sku}.{current_sku}.{amazonEC2_offer_code}.priceDimensions",
            "map_key",
        ):
            current_rate_code = value
        elif (prefix, event) == (
            f"terms.OnDemand.{current_sku}.{current_sku}.{amazonEC2_offer_code}.priceDimensions"
            f".{current_rate_code}.unit",
            "string",
        ):
            prices[current_sku]["unit"] = value
        elif (prefix, event) == (
            f"terms.OnDemand.{current_sku}.{current_sku}.{amazonEC2_offer_code}.priceDimensions"
            f".{current_rate_code}.pricePerUnit.USD",
            "string",
        ):
            prices[current_sku]["price"] = value
    return prices

# For each combination of location - size - os the file has a different sku.
# For each sku we have a price
def scrape_ec2_pricing():
    skus = {}
    prices = get_all_prices()
    json_file = get_json()
    parser = ijson.parse(json_file)
    current_sku = ""

    for prefix, event, value in parser:
        if "terms" in prefix:
            break
        if (prefix, event) == ("products", "map_key"):
            current_sku = value
            skus[current_sku] = {"sku": value}
        elif (prefix, event) == (f"products.{current_sku}.productFamily", "string"):
            skus[current_sku]["family"] = value
        elif (prefix, event) == (
            f"products.{current_sku}.attributes.location",
            "string",
        ):
            skus[current_sku]["locationName"] = value
        elif (prefix, event) == (
            f"products.{current_sku}.attributes.locationType",
            "string",
        ):
            skus[current_sku]["locationType"] = value
        elif (prefix, event) == (
            f"products.{current_sku}.attributes.instanceType",
            "string",
        ):
            skus[current_sku]["size"] = value
        elif (prefix, event) == (
            f"products.{current_sku}.attributes.operatingSystem",
            "string",
        ):
            skus[current_sku]["os"] = value
        elif (prefix, event) == (
            f"products.{current_sku}.attributes.regionCode",
            "string",
        ):
            skus[current_sku]["location"] = value
        # only get prices of compute instances atm
        elif (prefix, event) == (f"products.{current_sku}", "end_map"):
            if (
                "Compute Instance" not in skus[current_sku]["family"]
                and "Dedicated Host" not in skus[current_sku]["family"]
            ):
                del skus[current_sku]
    ec2_linux = defaultdict(OrderedDict)
    ec2_windows = defaultdict(OrderedDict)
    ec2_rhel = defaultdict(OrderedDict)
    ec2_rhel_ha = defaultdict(OrderedDict)
    ec2_suse = defaultdict(OrderedDict)

    os_map = {
        "Linux": ec2_linux,
        "Windows": ec2_windows,
        "RHEL": ec2_rhel,
        "SUSE": ec2_suse,
        "Red Hat Enterprise Linux with HA": ec2_rhel_ha,
    }
    for sku in skus:
        if skus[sku]["locationType"] != "AWS Region":
            continue
        os = skus[sku]["os"]
        if os == "NA":
            os = "Linux"  # linux is free
        os_dict = os_map.get(os)
        # new OS, until it is documented skip it
        if os_dict is None:
            print(f"Unexpected OS {os}")
            continue
        size = skus[sku]["size"]
        location = skus[sku]["location"]
        # size is first seen
        if not os_dict.get(size):
            os_dict[size] = {}
        # if price is not a number then label it as not available
        try:
            price = float(prices[sku]["price"])
            if os_dict[size].get(location) and os_dict[size][location] > price:
                # do nothing, keep the highest price
                pass
            else:
                os_dict[size][location] = price
        except ValueError:
            os_dict[size][location] = "n/a"
        except KeyError:
            # size is available only reserved
            del os_dict[size]
    return {
        "ec2_linux": ec2_linux,
        "ec2_windows": ec2_windows,
        "ec2_rhel": ec2_rhel,
        "ec2_suse": ec2_suse,
        "ec2_rhel_ha": ec2_rhel_ha,
    }

def update_pricing_file(pricing_file_path, pricing_data):
    with open(pricing_file_path, "r") as fp:
        content = fp.read()

    data = json.loads(content)
    original_data = copy.deepcopy(data)

    data["compute"].update(pricing_data)

    if data == original_data:
        # Nothing has changed, bail out early and don't update "updated" attribute
        print("Nothing has changed, skipping update.")
        return

    data["updated"] = int(time.time())

    # Always sort the pricing info
    data = sort_nested_dict(data)

    content = json.dumps(data, indent=4)
    lines = content.splitlines()
    lines = [line.rstrip() for line in lines]
    content = "\n".join(lines)

    with open(pricing_file_path, "w") as fp:
        fp.write(content)

def sort_nested_dict(value):
    """
    Recursively sort a nested dict.
    """
    result = OrderedDict()

    for key, value in sorted(value.items(), key=sort_key_by_numeric_other):
        if isinstance(value, (dict, OrderedDict)):
            result[key] = sort_nested_dict(value)
        else:
            result[key] = value

    return result

def sort_key_by_numeric_other(key_value):
    """
    Split key into numeric, alpha and other part and sort accordingly.
    """
    result = []

    for (numeric, alpha, other) in RE_NUMERIC_OTHER.findall(key_value[0]):
        numeric = int(numeric) if numeric else -1
        alpha = INSTANCE_SIZES.index(alpha) if alpha in INSTANCE_SIZES else alpha
        alpha = str(alpha)
        item = tuple([numeric, alpha, other])
        result.append(item)

    return tuple(result)

def main():
    print(
        "Scraping EC2 pricing data (if this runs for the first time "
        "it has to download a 3GB file, depending on your bandwith "
        "it might take a while)...."
    )

    pricing_data = scrape_ec2_pricing()
    update_pricing_file(pricing_file_path=PRICING_FILE_PATH, pricing_data=pricing_data)

    print("Pricing data updated")

if __name__ == "__main__":
    main()
<|endoftext|>