content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import json
  from typing import (
      Any, Dict, Tuple,
  )
  
  import dateutil.parser
  from pyhocon import ConfigFactory, ConfigTree
  
  from databuilder.extractor.base_extractor import Extractor
  from databuilder.models.table_stats import TableColumnStats
  
  
  class PandasProfilingColumnStatsExtractor(Extractor):
      FILE_PATH = 'file_path'
      DATABASE_NAME = 'database_name'
      TABLE_NAME = 'table_name'
      SCHEMA_NAME = 'schema_name'
      CLUSTER_NAME = 'cluster_name'
  
      # if you wish to collect only selected set of metrics configure stat_mappings option of the extractor providing
      # similar dictionary but containing only keys of metrics you wish to collect.
      # For example - if you want only min and max value of a column, provide extractor with configuration option:
      # PandasProfilingColumnStatsExtractor.STAT_MAPPINGS = {'max': ('Maximum', float), 'min': ('Minimum', float)}
      STAT_MAPPINGS = 'stat_mappings'
  
      # - key - raw name of the stat in pandas-profiling. Value - tuple of stat spec.
      # - first value of the tuple - full name of the stat
      # - second value of the tuple - function modifying the stat (by default we just do type casting)
      DEFAULT_STAT_MAPPINGS = {
          '25%': ('Quantile 25%', float),
          '5%': ('Quantile 5%', float),
          '50%': ('Quantile 50%', float),
          '75%': ('Quantile 75%', float),
          '95%': ('Quantile 95%', float),
          'chi_squared': ('Chi squared', lambda x: float(x.get('statistic'))),
          'count': ('Count', int),
          'is_unique': ('Unique', bool),
          'kurtosis': ('Kurtosis', float),
          'max': ('Maximum', str),
          'max_length': ('Maximum length', int),
          'mean': ('Mean', float),
          'mean_length': ('Mean length', int),
          'median_length': ('Median length', int),
          'min': ('Minimum', str),
          'min_length': ('Minimum length', int),
          'monotonic': ('Monotonic', bool),
          'n_characters': ('Characters', int),
          'n_characters_distinct': ('Distinct characters', int),
          'n_distinct': ('Distinct values', int),
          'n_infinite': ('Infinite values', int),
          'n_missing': ('Missing values', int),
          'n_negative': ('Negative values', int),
          'n_unique': ('Unique values', int),
          'n_zeros': ('Zeros', int),
          'p_distinct': ('Distinct values %', lambda x: float(x * 100)),
          'p_infinite': ('Infinite values %', lambda x: float(x * 100)),
          'p_missing': ('Missing values %', lambda x: float(x * 100)),
          'p_negative': ('Negative values %', lambda x: float(x * 100)),
          'p_unique': ('Unique values %', lambda x: float(x * 100)),
          'p_zeros': ('Zeros %', lambda x: float(x * 100)),
          'range': ('Range', str),
          'skewness': ('Skewness', float),
          'std': ('Std. deviation', float),
          'sum': ('Sum', float),
          'variance': ('Variance', float)
          # Stats available in pandas-profiling but are not collected by default and require custom, conscious config..
          # 'block_alias_char_counts': ('',),
          # 'block_alias_counts': ('',),
          # 'block_alias_values': ('',),
          # 'category_alias_char_counts': ('',),
          # 'category_alias_counts': ('',),
          # 'category_alias_values': ('',),
          # 'character_counts': ('',),
          # 'cv': ('',),
          # 'first_rows': ('',),
          # 'hashable': ('',),
          # 'histogram': ('',),
          # 'histogram_frequencies': ('',),
          # 'histogram_length': ('',),
          # 'iqr': ('',),
          # 'length': ('',),
          # 'mad': ('',),
          # 'memory_size': ('',),
          # 'monotonic_decrease': ('Monotonic decrease', bool),
          # 'monotonic_decrease_strict': ('Strict monotonic decrease', bool),
          # 'monotonic_increase': ('Monotonic increase', bool),
          # 'monotonic_increase_strict': ('Strict monotonic increase', bool),
          # 'n': ('',),
          # 'n_block_alias': ('',),
          # 'n_category': ('Categories', int),
          # 'n_scripts': ('',),
          # 'ordering': ('',),
          # 'script_char_counts': ('',),
          # 'script_counts': ('',),
          # 'value_counts_index_sorted': ('',),
          # 'value_counts_without_nan': ('',),
          # 'word_counts': ('',),
          # 'type': ('Type', str)
      }
  
      PRECISION = 'precision'
  
      DEFAULT_CONFIG = ConfigFactory.from_dict({STAT_MAPPINGS: DEFAULT_STAT_MAPPINGS, PRECISION: 3})
  
      def get_scope(self) -> str:
          return 'extractor.pandas_profiling'
  
      def init(self, conf: ConfigTree) -> None:
          self.conf = conf.with_fallback(PandasProfilingColumnStatsExtractor.DEFAULT_CONFIG)
  
          self._extract_iter = self._get_extract_iter()
  
      def extract(self) -> Any:
          try:
              result = next(self._extract_iter)
  
              return result
          except StopIteration:
              return None
  
      def _get_extract_iter(self) -> Any:
          report = self._load_report()
  
          variables = report.get('variables', dict())
          report_time = self.parse_date(report.get('analysis', dict()).get('date_start'))
  
          for column_name, column_stats in variables.items():
              for _stat_name, stat_value in column_stats.items():
                  stat_spec = self.stat_mappings.get(_stat_name)
  
                  if stat_spec:
                      stat_name, stat_modifier = stat_spec
  
                      if isinstance(stat_value, float):
                          stat_value = self.round_value(stat_value)
  
                      stat = TableColumnStats(table_name=self.table_name, col_name=column_name, stat_name=stat_name,
                                              stat_val=stat_modifier(stat_value), start_epoch=report_time, end_epoch='0',
                                              db=self.database_name, cluster=self.cluster_name, schema=self.schema_name)
  
                      yield stat
  
      def _load_report(self) -> Dict[str, Any]:
          path = self.conf.get(PandasProfilingColumnStatsExtractor.FILE_PATH)
  
          try:
              with open(path, 'r') as f:
                  _data = f.read()
  
              data = json.loads(_data)
  
              return data
          except Exception:
              return {}
  
      @staticmethod
      def parse_date(string_date: str) -> str:
          try:
              date_parsed = dateutil.parser.parse(string_date)
  
              # date from pandas-profiling doesn't contain timezone so to be timezone safe we need to assume it's utc
              if not date_parsed.tzname():
                  return PandasProfilingColumnStatsExtractor.parse_date(f'{string_date}+0000')
  
              return str(int(date_parsed.timestamp()))
          except Exception:
              return '0'
  
      def round_value(self, value: float) -> float:
          return round(value, self.conf.get(PandasProfilingColumnStatsExtractor.PRECISION))
  
      @property
      def stat_mappings(self) -> Dict[str, Tuple[str, Any]]:
          return dict(self.conf.get(PandasProfilingColumnStatsExtractor.STAT_MAPPINGS))
  
      @property
      def cluster_name(self) -> str:
          return self.conf.get(PandasProfilingColumnStatsExtractor.CLUSTER_NAME)
  
      @property
      def database_name(self) -> str:
          return self.conf.get(PandasProfilingColumnStatsExtractor.DATABASE_NAME)
  
      @property
      def schema_name(self) -> str:
          return self.conf.get(PandasProfilingColumnStatsExtractor.SCHEMA_NAME)
  
      @property
      def table_name(self) -> str:
          return self.conf.get(PandasProfilingColumnStatsExtractor.TABLE_NAME)
  