PyPI - upgini - Versions diffs - 1.1.280a3418.post2__py3-none-any.whl → 1.2.31a1__py3-none-any.whl - Mend

upgini 1.1.280a3418.post2py3-none-any.whl → 1.2.31a1py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.

This version of upgini might be problematic. Click here for more details.

Files changed (43) hide show

upgini/__about__.py +1 -1
upgini/__init__.py +4 -20
upgini/autofe/all_operands.py +39 -10
upgini/autofe/binary.py +148 -45
upgini/autofe/date.py +197 -26
upgini/autofe/feature.py +102 -19
upgini/autofe/groupby.py +22 -22
upgini/autofe/operand.py +9 -6
upgini/autofe/unary.py +78 -54
upgini/autofe/vector.py +8 -8
upgini/data_source/data_source_publisher.py +128 -5
upgini/dataset.py +50 -386
upgini/features_enricher.py +936 -541
upgini/http.py +27 -16
upgini/lazy_import.py +35 -0
upgini/metadata.py +84 -59
upgini/metrics.py +164 -34
upgini/normalizer/normalize_utils.py +197 -0
upgini/resource_bundle/strings.properties +66 -51
upgini/search_task.py +10 -4
upgini/utils/Roboto-Regular.ttf +0 -0
upgini/utils/base_search_key_detector.py +14 -12
upgini/utils/country_utils.py +16 -0
upgini/utils/custom_loss_utils.py +39 -36
upgini/utils/datetime_utils.py +98 -45
upgini/utils/deduplicate_utils.py +135 -112
upgini/utils/display_utils.py +46 -15
upgini/utils/email_utils.py +54 -16
upgini/utils/feature_info.py +172 -0
upgini/utils/features_validator.py +34 -20
upgini/utils/ip_utils.py +100 -1
upgini/utils/phone_utils.py +343 -0
upgini/utils/postal_code_utils.py +34 -0
upgini/utils/sklearn_ext.py +28 -19
upgini/utils/target_utils.py +113 -57
upgini/utils/warning_counter.py +1 -0
upgini/version_validator.py +8 -4
{upgini-1.1.280a3418.post2.dist-info → upgini-1.2.31a1.dist-info}/METADATA +31 -16
upgini-1.2.31a1.dist-info/RECORD +65 -0
upgini/normalizer/phone_normalizer.py +0 -340
upgini-1.1.280a3418.post2.dist-info/RECORD +0 -62
{upgini-1.1.280a3418.post2.dist-info → upgini-1.2.31a1.dist-info}/WHEEL +0 -0
{upgini-1.1.280a3418.post2.dist-info → upgini-1.2.31a1.dist-info}/licenses/LICENSE +0 -0

upgini/utils/display_utils.py CHANGED Viewed

@@ -4,12 +4,14 @@ import textwrap
 import urllib.parse
 import uuid
 from datetime import datetime, timezone
-from io import BytesIO
+from io import StringIO
 from typing import Callable, List, Optional
 import pandas as pd
 from xhtml2pdf import pisa
+from upgini.__about__ import __version__
 def ipython_available() -> bool:
     try:
@@ -71,7 +73,9 @@ def make_table(df: pd.DataFrame, wrap_long_string=None) -> str:
     )
-def display_html_dataframe(df: pd.DataFrame, internal_df: pd.DataFrame, header: str):
+def display_html_dataframe(
+    df: pd.DataFrame, internal_df: pd.DataFrame, header: str, display_id: Optional[str] = None, display_handle=None
+):
     if not ipython_available():
         print(header)
         print(internal_df)
@@ -132,7 +136,10 @@ def display_html_dataframe(df: pd.DataFrame, internal_df: pd.DataFrame, header:
             {table_html}
         </div>
         """
-    display(HTML(result_html))
+    if display_handle:
+        return display_handle.update(HTML(result_html))
+    else:
+        return display(HTML(result_html), display_id=display_id)
 def make_html_report(
@@ -143,7 +150,7 @@ def make_html_report(
     search_id: str,
     email: Optional[str] = None,
     search_keys: Optional[List[str]] = None,
-):
+) -> str:
     # relevant_features_df = relevant_features_df.copy()
     # relevant_features_df["Feature name"] = relevant_features_df["Feature name"].apply(
     #     lambda x: "*" + x if x.contains("_autofe_") else x
@@ -154,9 +161,18 @@ def make_html_report(
         """<button type="button">Request a quote</button></a>"""
     )
     relevant_datasources_df.rename(columns={"action": "&nbsp;"}, inplace=True)
+    try:
+        from importlib.resources import files
+        font_path = files('upgini.utils').joinpath('Roboto-Regular.ttf')
+    except Exception:
+        from pkg_resources import resource_filename
+        font_path = resource_filename('upgini.utils', 'Roboto-Regular.ttf')
     return f"""<html>
         <head>
             <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
+            <meta charset="UTF-8">
             <style>
                 @page {{
                     size: a4 portrait;
@@ -166,23 +182,25 @@ def make_html_report(
                         /*-pdf-frame-border: 1;*/
                     }}
                     @frame content_frame {{
-                        left: 10pt; width: 574pt; top: 50pt; height: 752pt;
+                        left: 10pt; width: 574pt; top: 50pt; height: 742pt;
                         /*-pdf-frame-border: 1;*/
                     }}
                     @frame footer_frame {{
                         -pdf-frame-content: footer_content;
-                        left: 10pt; width: 574pt; top: 802pt; height: 30pt;
+                        left: 10pt; width: 574pt; top: 802pt; height: 40pt;
                         /*-pdf-frame-border: 1;*/
                     }}
                 }}
                 @font-face {{
-                    font-family: "Alice-Regular";
-                    src: url("/fonts/Alice-Regular.ttf") format("truetype");
+                    font-family: "Roboto";
+                    src: url("{font_path}") format("truetype");
                 }}
                 body {{
-                    font-family: "Alice-Regular", Arial, sans-serif;
+                    font-family: "Roboto", sans-serif;
+                    font-weight: 400;
+                    font-style: normal;
                 }}
                 #header_content {{
@@ -234,7 +252,8 @@ def make_html_report(
             <div id="header_content">UPGINI</div>
             <div id="footer_content">
                 © Upgini</br>
-                sales@upgini.com
+                sales@upgini.com</br>
+                Launched by version {__version__}
             </div>
             <h1>Data search report</h1>
@@ -257,7 +276,7 @@ def make_html_report(
             }
             <h3>Relevant data sources</h3>
             {make_table(relevant_datasources_df)}
-            <h3>All relevant features. Listing</h3>
+            <h3>All relevant features. Listing ({len(relevant_features_df)} items)</h3>
             {make_table(relevant_features_df, wrap_long_string=25)}
             {"<h3>Description of AutoFE feature names</h3>" + make_table(autofe_descriptions_df, wrap_long_string=25)
              if autofe_descriptions_df is not None
@@ -277,6 +296,8 @@ def prepare_and_show_report(
     search_id: str,
     email: Optional[str],
     search_keys: Optional[List[str]] = None,
+    display_id: Optional[str] = None,
+    display_handle=None,
 ):
     if not ipython_available():
         return
@@ -286,22 +307,32 @@ def prepare_and_show_report(
     )
     if len(relevant_features_df) > 0:
-        show_button_download_pdf(report)
+        return show_button_download_pdf(report, display_id=display_id, display_handle=display_handle)
-def show_button_download_pdf(source: str, title="\U0001F4CA Download PDF report"):
+def show_button_download_pdf(
+    source: str, title="\U0001F4CA Download PDF report", display_id: Optional[str] = None, display_handle=None
+):
     from IPython.display import HTML, display
     file_name = f"upgini-report-{uuid.uuid4()}.pdf"
+    # from weasyprint import HTML
+    # html = HTML(string=source)
+    # html.write_pdf(file_name)
     with open(file_name, "wb") as output:
-        pisa.CreatePDF(src=BytesIO(source.encode("UTF-8")), dest=output)
+        pisa.CreatePDF(src=StringIO(source), dest=output, encoding="UTF-8")
     with open(file_name, "rb") as f:
         b64 = base64.b64encode(f.read())
         payload = b64.decode()
         html = f"""<a download="{file_name}" href="data:application/pdf;base64,{payload}" target="_blank">
         <button>{title}</button></a>"""
-        display(HTML(html))
+        if display_handle is not None:
+            display_handle.update(HTML(html))
+        else:
+            return display(HTML(html), display_id=display_id)
 def show_request_quote_button():

upgini/utils/email_utils.py CHANGED Viewed

@@ -7,7 +7,7 @@ import pandas as pd
 from pandas.api.types import is_object_dtype, is_string_dtype
 from upgini.metadata import SearchKey
-from upgini.resource_bundle import bundle
+from upgini.resource_bundle import ResourceBundle, get_custom_bundle
 from upgini.utils.base_search_key_detector import BaseSearchKeyDetector
 EMAIL_REGEX = re.compile(r"^[a-zA-Z0-9.!#$%&’*+/=?^_`{|}~-]+@[a-zA-Z0-9-]+(?:\.[a-zA-Z0-9-]+)*$")
@@ -28,27 +28,54 @@ class EmailSearchKeyDetector(BaseSearchKeyDetector):
         return is_email_count / all_count > 0.1
+class EmailDomainGenerator:
+    DOMAIN_SUFFIX = "_domain"
+    def __init__(self, email_columns: List[str]):
+        self.email_columns = email_columns
+        self.generated_features = []
+    def generate(self, df: pd.DataFrame) -> pd.DataFrame:
+        for email_col in self.email_columns:
+            domain_feature = email_col + self.DOMAIN_SUFFIX
+            if domain_feature not in df.columns:
+                df[domain_feature] = df[email_col].apply(self._email_to_domain).astype("string")
+                self.generated_features.append(domain_feature)
+        return df
+    @staticmethod
+    def _email_to_domain(email: str) -> Optional[str]:
+        if email is not None and isinstance(email, str) and "@" in email:
+            name_and_domain = email.split("@")
+            if len(name_and_domain) == 2 and len(name_and_domain[1]) > 0:
+                return name_and_domain[1]
 class EmailSearchKeyConverter:
-    HEM_COLUMN_NAME = "hashed_email"
-    DOMAIN_COLUMN_NAME = "email_domain"
-    EMAIL_ONE_DOMAIN_COLUMN_NAME = "email_one_domain"
+    HEM_SUFFIX = "_hem"
+    ONE_DOMAIN_SUFFIX = "_one_domain"
     def __init__(
         self,
         email_column: str,
         hem_column: Optional[str],
         search_keys: Dict[str, SearchKey],
+        columns_renaming: Dict[str, str],
+        unnest_search_keys: Optional[List[str]] = None,
+        bundle: Optional[ResourceBundle] = None,
         logger: Optional[logging.Logger] = None,
     ):
         self.email_column = email_column
         self.hem_column = hem_column
         self.search_keys = search_keys
+        self.columns_renaming = columns_renaming
+        self.unnest_search_keys = unnest_search_keys
+        self.bundle = bundle or get_custom_bundle()
         if logger is not None:
             self.logger = logger
         else:
             self.logger = logging.getLogger()
             self.logger.setLevel("FATAL")
-        self.generated_features: List[str] = []
         self.email_converted_to_hem = False
     @staticmethod
@@ -59,7 +86,7 @@ class EmailSearchKeyConverter:
         if not EMAIL_REGEX.fullmatch(email):
             return None
-        return sha256(email.lower().encode("utf-8")).hexdigest()
+        return sha256(email.lower().encode("utf-8")).hexdigest().lower()
     @staticmethod
     def _email_to_one_domain(email: str) -> Optional[str]:
@@ -70,25 +97,36 @@ class EmailSearchKeyConverter:
     def convert(self, df: pd.DataFrame) -> pd.DataFrame:
         df = df.copy()
+        original_email_column = self.columns_renaming[self.email_column]
         if self.hem_column is None:
-            df[self.HEM_COLUMN_NAME] = df[self.email_column].apply(self._email_to_hem)
-            if df[self.HEM_COLUMN_NAME].isna().all():
-                msg = bundle.get("all_emails_invalid").format(self.email_column)
+            hem_name = self.email_column + self.HEM_SUFFIX
+            df[hem_name] = df[self.email_column].apply(self._email_to_hem)
+            if df[hem_name].isna().all():
+                msg = self.bundle.get("all_emails_invalid").format(self.email_column)
                 print(msg)
                 self.logger.warning(msg)
-                df = df.drop(columns=self.HEM_COLUMN_NAME)
+                df = df.drop(columns=hem_name)
                 del self.search_keys[self.email_column]
                 return df
-            self.search_keys[self.HEM_COLUMN_NAME] = SearchKey.HEM
+            self.search_keys[hem_name] = SearchKey.HEM
+            if self.email_column in self.unnest_search_keys:
+                self.unnest_search_keys.append(hem_name)
+            self.columns_renaming[hem_name] = original_email_column  # it could be upgini_email_unnest...
             self.email_converted_to_hem = True
+        else:
+            df[self.hem_column] = df[self.hem_column].astype("string").str.lower()
         del self.search_keys[self.email_column]
+        if self.email_column in self.unnest_search_keys:
+            self.unnest_search_keys.remove(self.email_column)
-        df[self.EMAIL_ONE_DOMAIN_COLUMN_NAME] = df[self.email_column].apply(self._email_to_one_domain)
-        self.search_keys[self.EMAIL_ONE_DOMAIN_COLUMN_NAME] = SearchKey.EMAIL_ONE_DOMAIN
+        one_domain_name = self.email_column + self.ONE_DOMAIN_SUFFIX
+        df[one_domain_name] = df[self.email_column].apply(self._email_to_one_domain)
+        self.columns_renaming[one_domain_name] = original_email_column
+        self.search_keys[one_domain_name] = SearchKey.EMAIL_ONE_DOMAIN
-        df[self.DOMAIN_COLUMN_NAME] = df[self.EMAIL_ONE_DOMAIN_COLUMN_NAME].str[1:]
-        self.generated_features.append(self.DOMAIN_COLUMN_NAME)
+        if self.email_converted_to_hem:
+            df = df.drop(columns=self.email_column)
+            del self.columns_renaming[self.email_column]
         return df

upgini/utils/feature_info.py ADDED Viewed

@@ -0,0 +1,172 @@
+from dataclasses import dataclass
+import itertools
+from typing import Dict, List
+import numpy as np
+import pandas as pd
+from upgini.metadata import FeaturesMetadataV2
+from upgini.resource_bundle import ResourceBundle
+LLM_SOURCE = "LLM with external data augmentation"
+@dataclass
+class FeatureInfo:
+    name: str
+    internal_name: str
+    rounded_shap: float
+    hitrate: float
+    value_preview: str
+    provider: str
+    internal_provider: str
+    source: str
+    internal_source: str
+    update_frequency: str
+    commercial_schema: str
+    doc_link: str
+    data_provider_link: str
+    data_source_link: str
+    @staticmethod
+    def from_metadata(feature_meta: FeaturesMetadataV2, data: pd.DataFrame, is_client_feature: bool) -> "FeatureInfo":
+        return FeatureInfo(
+            name=_get_name(feature_meta),
+            internal_name=_get_internal_name(feature_meta),
+            rounded_shap=_round_shap_value(feature_meta.shap_value),
+            hitrate=feature_meta.hit_rate,
+            value_preview=_get_feature_sample(feature_meta, data),
+            provider=_get_provider(feature_meta, is_client_feature),
+            internal_provider=_get_internal_provider(feature_meta, is_client_feature),
+            source=_get_source(feature_meta, is_client_feature),
+            internal_source=_get_internal_source(feature_meta, is_client_feature),
+            update_frequency=feature_meta.update_frequency,
+            commercial_schema=feature_meta.commercial_schema,
+            doc_link=feature_meta.doc_link,
+            data_provider_link=feature_meta.data_provider_link,
+            data_source_link=feature_meta.data_source_link,
+        )
+    def to_row(self, bundle: ResourceBundle) -> Dict[str, str]:
+        return {
+            bundle.get("features_info_name"): self.name,
+            bundle.get("features_info_shap"): self.rounded_shap,
+            bundle.get("features_info_hitrate"): self.hitrate,
+            bundle.get("features_info_value_preview"): self.value_preview,
+            bundle.get("features_info_provider"): self.provider,
+            bundle.get("features_info_source"): self.source,
+            bundle.get("features_info_update_frequency"): self.update_frequency,
+        }
+    def to_row_without_links(self, bundle: ResourceBundle) -> Dict[str, str]:
+        return {
+            bundle.get("features_info_name"): self.internal_name,
+            bundle.get("features_info_shap"): self.rounded_shap,
+            bundle.get("features_info_hitrate"): self.hitrate,
+            bundle.get("features_info_value_preview"): self.value_preview,
+            bundle.get("features_info_provider"): self.internal_provider,
+            bundle.get("features_info_source"): self.internal_source,
+            bundle.get("features_info_update_frequency"): self.update_frequency,
+        }
+    def to_internal_row(self, bundle: ResourceBundle) -> Dict[str, str]:
+        return {
+            bundle.get("features_info_name"): self.internal_name,
+            "feature_link": self.doc_link,
+            bundle.get("features_info_shap"): self.rounded_shap,
+            bundle.get("features_info_hitrate"): self.hitrate,
+            bundle.get("features_info_value_preview"): self.value_preview,
+            bundle.get("features_info_provider"): self.internal_provider,
+            "provider_link": self.data_provider_link,
+            bundle.get("features_info_source"): self.internal_source,
+            "source_link": self.data_source_link,
+            bundle.get("features_info_commercial_schema"): self.commercial_schema or "",
+            bundle.get("features_info_update_frequency"): self.update_frequency,
+        }
+def _get_feature_sample(feature_meta: FeaturesMetadataV2, data: pd.DataFrame) -> str:
+    if feature_meta.name in data.columns:
+        feature_sample = np.random.choice(data[feature_meta.name].dropna().unique(), 3).tolist()
+        if len(feature_sample) > 0 and isinstance(feature_sample[0], float):
+            feature_sample = [round(f, 4) for f in feature_sample]
+        feature_sample = [str(f) for f in feature_sample]
+        feature_sample = ", ".join(feature_sample)
+        if len(feature_sample) > 30:
+            feature_sample = feature_sample[:30] + "..."
+    else:
+        feature_sample = ""
+    return feature_sample
+def _get_name(feature_meta: FeaturesMetadataV2) -> str:
+    if feature_meta.doc_link:
+        return _to_anchor(feature_meta.doc_link, feature_meta.name)
+    else:
+        return feature_meta.name
+def _get_internal_name(feature_meta: FeaturesMetadataV2) -> str:
+    return feature_meta.name
+def _get_provider(feature_meta: FeaturesMetadataV2, is_client_feature: bool) -> str:
+    providers = _list_or_single(feature_meta.data_providers, feature_meta.data_provider)
+    provider_links = _list_or_single(feature_meta.data_provider_links, feature_meta.data_provider_link)
+    if providers:
+        provider = _make_links(providers, provider_links)
+    else:
+        provider = "" if is_client_feature else _to_anchor("https://upgini.com", "Upgini")
+    return provider
+def _get_internal_provider(feature_meta: FeaturesMetadataV2, is_client_feature: bool) -> str:
+    return "" if is_client_feature else (feature_meta.data_provider or "Upgini")
+def _get_source(feature_meta: FeaturesMetadataV2, is_client_feature: bool) -> str:
+    sources = _list_or_single(feature_meta.data_sources, feature_meta.data_source)
+    source_links = _list_or_single(feature_meta.data_source_links, feature_meta.data_source_link)
+    if sources:
+        source = _make_links(sources, source_links)
+    else:
+        source = _get_internal_source(feature_meta, is_client_feature)
+    return source
+def _get_internal_source(feature_meta: FeaturesMetadataV2, is_client_feature: bool) -> str:
+    return feature_meta.data_source or (
+        LLM_SOURCE
+        if not feature_meta.name.endswith("_country")
+        and not feature_meta.name.endswith("_postal_code")
+        and not is_client_feature
+        else ""
+    )
+def _list_or_single(lst: List[str], single: str):
+    return lst or ([single] if single else [])
+def _to_anchor(link: str, value: str) -> str:
+    if not value:
+        return ""
+    elif not link:
+        return value
+    elif value == LLM_SOURCE:
+        return value
+    else:
+        return f"<a href='{link}' target='_blank' rel='noopener noreferrer'>{value}</a>"
+def _make_links(names: List[str], links: List[str]):
+    all_links = [_to_anchor(link, name) for name, link in itertools.zip_longest(names, links)]
+    return ",".join(all_links)
+def _round_shap_value(shap: float) -> float:
+    if shap > 0.0 and shap < 0.0001:
+        return 0.0001
+    else:
+        return round(shap, 4)

upgini/utils/features_validator.py CHANGED Viewed

@@ -1,12 +1,12 @@
 import logging
 from logging import Logger
-from typing import List, Optional
+from typing import Dict, List, Optional, Tuple
+import numpy as np
 import pandas as pd
 from pandas.api.types import is_integer_dtype, is_object_dtype, is_string_dtype
 from upgini.resource_bundle import bundle
-from upgini.utils.warning_counter import WarningCounter
 class FeaturesValidator:
@@ -21,12 +21,13 @@ class FeaturesValidator:
         self,
         df: pd.DataFrame,
         features: List[str],
-        features_for_generate: Optional[List[str]],
-        warning_counter: WarningCounter,
-    ) -> List[str]:
+        features_for_generate: Optional[List[str]] = None,
+        columns_renaming: Optional[Dict[str, str]] = None,
+    ) -> Tuple[List[str], List[str]]:
         # one_hot_encoded_features = []
         empty_or_constant_features = []
         high_cardinality_features = []
+        warnings = []
         for f in features:
             column = df[f]
@@ -51,26 +52,28 @@ class FeaturesValidator:
         # if one_hot_encoded_features:
         #     msg = bundle.get("one_hot_encoded_features").format(one_hot_encoded_features)
-        #     print(msg)
-        #     self.logger.warning(msg)
-        #     warning_counter.increment()
+        #     warnings.append(msg)
+        columns_renaming = columns_renaming or {}
         if empty_or_constant_features:
-            msg = bundle.get("empty_or_contant_features").format(empty_or_constant_features)
-            print(msg)
-            self.logger.warning(msg)
-            warning_counter.increment()
+            msg = bundle.get("empty_or_contant_features").format(
+                [columns_renaming.get(f, f) for f in empty_or_constant_features]
+            )
+            warnings.append(msg)
         high_cardinality_features = self.find_high_cardinality(df[features])
         if features_for_generate:
-            high_cardinality_features = [f for f in high_cardinality_features if f not in features_for_generate]
+            high_cardinality_features = [
+                f for f in high_cardinality_features if columns_renaming.get(f, f) not in features_for_generate
+            ]
         if high_cardinality_features:
-            msg = bundle.get("high_cardinality_features").format(high_cardinality_features)
-            print(msg)
-            self.logger.warning(msg)
-            warning_counter.increment()
+            msg = bundle.get("high_cardinality_features").format(
+                [columns_renaming.get(f, f) for f in high_cardinality_features]
+            )
+            warnings.append(msg)
-        return empty_or_constant_features + high_cardinality_features
+        return (empty_or_constant_features + high_cardinality_features, warnings)
     @staticmethod
     def find_high_cardinality(df: pd.DataFrame) -> List[str]:
@@ -81,10 +84,21 @@ class FeaturesValidator:
         return [
             i
             for i in df
-            if (is_object_dtype(df[i]) or is_string_dtype(df[i]) or is_integer_dtype(df[i]))
+            if (is_object_dtype(df[i]) or is_string_dtype(df[i]) or FeaturesValidator.__is_integer(df[i]))
             and (df[i].nunique(dropna=False) / row_count >= 0.85)
         ]
+    @staticmethod
+    def __is_integer(series: pd.Series) -> bool:
+        return (
+            is_integer_dtype(series)
+            or series.dropna()
+            .apply(
+                lambda f: (float.is_integer(f) and abs(f) < np.iinfo(np.int64).max) if isinstance(f, float) else False
+            )
+            .all()
+        )
     @staticmethod
     def find_constant_features(df: pd.DataFrame) -> List[str]:
-        return [i for i in df if df[i].nunique() == 1]
+        return [i for i in df if df[i].nunique() <= 1]

upgini/utils/ip_utils.py CHANGED Viewed

@@ -1,15 +1,114 @@
 import logging
-from typing import Dict, List, Optional
+from ipaddress import IPv4Address, IPv6Address, _BaseAddress, ip_address
+from typing import Dict, List, Optional, Union
 import pandas as pd
 from requests import get
+from upgini.errors import ValidationError
 from upgini.metadata import SearchKey
+from upgini.resource_bundle import ResourceBundle, get_custom_bundle
 # from upgini.resource_bundle import bundle
 # from upgini.utils.track_info import get_track_metrics
+class IpSearchKeyConverter:
+    def __init__(
+        self,
+        ip_column: str,
+        search_keys: Dict[str, SearchKey],
+        columns_renaming: Dict[str, str],
+        unnest_search_keys: Optional[List[str]] = None,
+        bundle: Optional[ResourceBundle] = None,
+        logger: Optional[logging.Logger] = None,
+    ):
+        self.ip_column = ip_column
+        self.search_keys = search_keys
+        self.columns_renaming = columns_renaming
+        self.unnest_search_keys = unnest_search_keys
+        self.bundle = bundle or get_custom_bundle()
+        if logger is not None:
+            self.logger = logger
+        else:
+            self.logger = logging.getLogger()
+            self.logger.setLevel("FATAL")
+    @staticmethod
+    def _ip_to_int(ip: Optional[_BaseAddress]) -> Optional[int]:
+        try:
+            if isinstance(ip, (IPv4Address, IPv6Address)):
+                return int(ip)
+        except Exception:
+            pass
+    @staticmethod
+    def _ip_to_int_str(ip: Optional[_BaseAddress]) -> Optional[str]:
+        try:
+            if isinstance(ip, (IPv4Address, IPv6Address)):
+                return str(int(ip))
+        except Exception:
+            pass
+    @staticmethod
+    def _safe_ip_parse(ip: Union[str, int, IPv4Address, IPv6Address]) -> Optional[_BaseAddress]:
+        try:
+            return ip_address(ip)
+        except ValueError:
+            pass
+    # @staticmethod
+    # def _is_ipv4(ip: Optional[_BaseAddress]):
+    #     return ip is not None and (
+    #         isinstance(ip, IPv4Address) or (isinstance(ip, IPv6Address) and ip.ipv4_mapped is not None)
+    #     )
+    # @staticmethod
+    # def _to_ipv4(ip: Optional[_BaseAddress]) -> Optional[IPv4Address]:
+    #     if isinstance(ip, IPv4Address):
+    #         return ip
+    #     return None
+    @staticmethod
+    def _to_ipv6(ip: Optional[_BaseAddress]) -> Optional[IPv6Address]:
+        if isinstance(ip, IPv6Address):
+            return ip
+        if isinstance(ip, IPv4Address):
+            return IPv6Address("::ffff:" + str(ip))
+        return None
+    def convert(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Convert ip address to int"""
+        self.logger.info("Convert ip address to int")
+        original_ip = self.columns_renaming[self.ip_column]
+        df[self.ip_column] = df[self.ip_column].apply(self._safe_ip_parse)
+        if df[self.ip_column].isnull().all():
+            raise ValidationError(self.bundle.get("invalid_ip").format(self.ip_column))
+        # legacy support
+        # ipv4 = self.ip_column + "_v4"
+        # df[ipv4] = df[self.ip_column].apply(self._to_ipv4).apply(self._ip_to_int).astype("Int64")
+        # self.search_keys[ipv4] = SearchKey.IP
+        # self.columns_renaming[ipv4] = original_ip
+        ipv6 = self.ip_column + "_v6"
+        df[ipv6] = (
+            df[self.ip_column]
+            .apply(self._to_ipv6)
+            .apply(self._ip_to_int_str)
+            .astype("string")
+            # .str.replace(".0", "", regex=False)
+        )
+        df = df.drop(columns=self.ip_column)
+        del self.search_keys[self.ip_column]
+        del self.columns_renaming[self.ip_column]
+        self.search_keys[ipv6] = SearchKey.IPV6_ADDRESS
+        self.columns_renaming[ipv6] = original_ip  # could be __unnest_ip...
+        return df
 class IpToCountrySearchKeyConverter:
     url = "http://ip-api.com/json/{}"

upgini 1.1.280a3418.post2__py3-none-any.whl → 1.2.31a1__py3-none-any.whl

Potentially problematic release.

upgini 1.1.280a3418.post2py3-none-any.whl → 1.2.31a1py3-none-any.whl