From f4e4bc51c8d0927867ea73a4fbc24efed71886fc Mon Sep 17 00:00:00 2001
From: Andrei Voinea <34933567+vandrw@users.noreply.github.com>
Date: Sat, 7 Sep 2024 12:27:38 +0000
Subject: [PATCH] Update deprecated code (#225)

* Update `pkg_resources`

* Update `pydantic` to match v2 requirements

* Bump numpy version to allow >2.0

* Update `typing` to `beartype.typing`
---
 alphafold3_pytorch/alphafold3.py                  |  9 ++++++++-
 alphafold3_pytorch/attention.py                   |  2 +-
 alphafold3_pytorch/common/amino_acid_constants.py |  2 +-
 alphafold3_pytorch/common/biomolecule.py          |  2 +-
 alphafold3_pytorch/common/dna_constants.py        |  2 +-
 alphafold3_pytorch/common/ligand_constants.py     |  2 +-
 alphafold3_pytorch/common/mmcif_metadata.py       |  2 +-
 alphafold3_pytorch/common/rna_constants.py        |  2 +-
 alphafold3_pytorch/configs.py                     | 15 ++++++++++-----
 alphafold3_pytorch/data/data_pipeline.py          |  2 +-
 alphafold3_pytorch/data/kalign.py                 |  2 +-
 alphafold3_pytorch/data/mmcif_parsing.py          |  2 +-
 alphafold3_pytorch/data/msa_parsing.py            |  2 +-
 alphafold3_pytorch/data/template_parsing.py       |  2 +-
 alphafold3_pytorch/data/weighted_pdb_sampler.py   |  2 +-
 alphafold3_pytorch/inputs.py                      | 11 ++++++++++-
 alphafold3_pytorch/life.py                        |  2 +-
 alphafold3_pytorch/trainer.py                     |  2 +-
 alphafold3_pytorch/utils/data_utils.py            |  2 +-
 alphafold3_pytorch/utils/model_utils.py           |  9 +++++----
 alphafold3_pytorch/utils/utils.py                 |  2 +-
 pyproject.toml                                    |  2 +-
 scripts/cluster_pdb_test_mmcifs.py                |  2 +-
 scripts/cluster_pdb_train_mmcifs.py               |  2 +-
 scripts/cluster_pdb_val_mmcifs.py                 |  2 +-
 25 files changed, 54 insertions(+), 32 deletions(-)

diff --git a/alphafold3_pytorch/alphafold3.py b/alphafold3_pytorch/alphafold3.py
index fe9a1837..ce873543 100644
--- a/alphafold3_pytorch/alphafold3.py
+++ b/alphafold3_pytorch/alphafold3.py
@@ -21,7 +21,14 @@
     Sequential,
 )
 
-from typing import Callable, Dict, List, Literal, NamedTuple, Tuple
+from beartype.typing import (
+    Callable,
+    Dict,
+    List,
+    Literal,
+    NamedTuple,
+    Tuple,
+)
 
 from alphafold3_pytorch.tensor_typing import (
     Float,
diff --git a/alphafold3_pytorch/attention.py b/alphafold3_pytorch/attention.py
index b5ccbcbd..6648bafd 100644
--- a/alphafold3_pytorch/attention.py
+++ b/alphafold3_pytorch/attention.py
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import NamedTuple, Tuple
+from beartype.typing import NamedTuple, Tuple
 
 import torch
 from torch import nn, Tensor
diff --git a/alphafold3_pytorch/common/amino_acid_constants.py b/alphafold3_pytorch/common/amino_acid_constants.py
index 40fbebd7..5f92b02e 100644
--- a/alphafold3_pytorch/common/amino_acid_constants.py
+++ b/alphafold3_pytorch/common/amino_acid_constants.py
@@ -1,6 +1,6 @@
 """Amino acid constants used in AlphaFold."""
 
-from typing import Final
+from beartype.typing import Final
 
 import numpy as np
 
diff --git a/alphafold3_pytorch/common/biomolecule.py b/alphafold3_pytorch/common/biomolecule.py
index ead95bec..c9303343 100644
--- a/alphafold3_pytorch/common/biomolecule.py
+++ b/alphafold3_pytorch/common/biomolecule.py
@@ -7,7 +7,7 @@
 import random
 from functools import partial
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Set, Tuple
+from beartype.typing import Any, Dict, List, Optional, Set, Tuple
 
 import numpy as np
 from Bio.PDB.mmcifio import MMCIFIO
diff --git a/alphafold3_pytorch/common/dna_constants.py b/alphafold3_pytorch/common/dna_constants.py
index 284f5455..a427ff9d 100644
--- a/alphafold3_pytorch/common/dna_constants.py
+++ b/alphafold3_pytorch/common/dna_constants.py
@@ -1,6 +1,6 @@
 """Deoxyribonucleic acid (DNA) constants used in AlphaFold."""
 
-from typing import Final
+from beartype.typing import Final
 
 import numpy as np
 
diff --git a/alphafold3_pytorch/common/ligand_constants.py b/alphafold3_pytorch/common/ligand_constants.py
index ac5a0bae..2a1840b3 100644
--- a/alphafold3_pytorch/common/ligand_constants.py
+++ b/alphafold3_pytorch/common/ligand_constants.py
@@ -1,6 +1,6 @@
 """Ligand constants used in AlphaFold."""
 
-from typing import Final
+from beartype.typing import Final
 
 import numpy as np
 
diff --git a/alphafold3_pytorch/common/mmcif_metadata.py b/alphafold3_pytorch/common/mmcif_metadata.py
index 6c613dfb..964a56b0 100644
--- a/alphafold3_pytorch/common/mmcif_metadata.py
+++ b/alphafold3_pytorch/common/mmcif_metadata.py
@@ -1,6 +1,6 @@
 """mmCIF metadata."""
 
-from typing import Mapping, Sequence
+from beartype.typing import Mapping, Sequence
 
 import numpy as np
 
diff --git a/alphafold3_pytorch/common/rna_constants.py b/alphafold3_pytorch/common/rna_constants.py
index 7172d8f3..c9238c8a 100644
--- a/alphafold3_pytorch/common/rna_constants.py
+++ b/alphafold3_pytorch/common/rna_constants.py
@@ -1,6 +1,6 @@
 """Ribonucleic acid (RNA) constants used in AlphaFold."""
 
-from typing import Final
+from beartype.typing import Final
 
 import numpy as np
 
diff --git a/alphafold3_pytorch/configs.py b/alphafold3_pytorch/configs.py
index 841290d4..bd3460ed 100644
--- a/alphafold3_pytorch/configs.py
+++ b/alphafold3_pytorch/configs.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from alphafold3_pytorch.tensor_typing import typecheck
-from typing import Callable, List, Dict, Literal
+from beartype.typing import Callable, List, Dict, Literal
 
 from alphafold3_pytorch.alphafold3 import Alphafold3
 
@@ -24,7 +24,11 @@
 import yaml
 from pathlib import Path
 
-from pydantic import BaseModel, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    model_validator,
+)
 
 from pydantic.types import (
     FilePath,
@@ -79,9 +83,10 @@ def yaml_config_path_to_dict(
 # base pydantic classes for constructing alphafold3 and trainer from config files
 
 class BaseModelWithExtra(BaseModel):
-    class Config:
-        extra = 'allow'
-        use_enum_values = True
+    model_config = ConfigDict(
+        extra = 'allow',
+        use_enum_values = True,
+    )
 
 class Alphafold3Config(BaseModelWithExtra):
     dim_atom_inputs: int
diff --git a/alphafold3_pytorch/data/data_pipeline.py b/alphafold3_pytorch/data/data_pipeline.py
index 54b19449..513a30ad 100644
--- a/alphafold3_pytorch/data/data_pipeline.py
+++ b/alphafold3_pytorch/data/data_pipeline.py
@@ -1,7 +1,7 @@
 """General-purpose data pipeline."""
 
 import os
-from typing import Dict, List, MutableMapping, Optional, Tuple
+from beartype.typing import Dict, List, MutableMapping, Optional, Tuple
 
 import numpy as np
 import torch
diff --git a/alphafold3_pytorch/data/kalign.py b/alphafold3_pytorch/data/kalign.py
index 4bb6b046..ab747ba9 100644
--- a/alphafold3_pytorch/data/kalign.py
+++ b/alphafold3_pytorch/data/kalign.py
@@ -17,7 +17,7 @@
 import subprocess  # nosec
 import tempfile
 from loguru import logger
-from typing import Mapping, Sequence, Tuple
+from beartype.typing import Mapping, Sequence, Tuple
 
 from alphafold3_pytorch.data import msa_parsing
 from alphafold3_pytorch.data.template_parsing import (
diff --git a/alphafold3_pytorch/data/mmcif_parsing.py b/alphafold3_pytorch/data/mmcif_parsing.py
index a7c15643..15588be5 100644
--- a/alphafold3_pytorch/data/mmcif_parsing.py
+++ b/alphafold3_pytorch/data/mmcif_parsing.py
@@ -7,7 +7,7 @@
 import logging
 from collections import defaultdict
 from operator import itemgetter
-from typing import Any, Mapping, Optional, Sequence, Set, Tuple
+from beartype.typing import Any, Mapping, Optional, Sequence, Set, Tuple
 
 import numpy as np
 from Bio import PDB
diff --git a/alphafold3_pytorch/data/msa_parsing.py b/alphafold3_pytorch/data/msa_parsing.py
index 01553fff..d23f2c37 100644
--- a/alphafold3_pytorch/data/msa_parsing.py
+++ b/alphafold3_pytorch/data/msa_parsing.py
@@ -6,7 +6,7 @@
 import random
 import re
 import string
-from typing import Literal, Optional, Sequence, Tuple
+from beartype.typing import Literal, Optional, Sequence, Tuple
 
 from alphafold3_pytorch.tensor_typing import typecheck
 
diff --git a/alphafold3_pytorch/data/template_parsing.py b/alphafold3_pytorch/data/template_parsing.py
index 72ff2586..7e9fd607 100644
--- a/alphafold3_pytorch/data/template_parsing.py
+++ b/alphafold3_pytorch/data/template_parsing.py
@@ -1,7 +1,7 @@
 import os
 from datetime import datetime
 from loguru import logger
-from typing import Any, Dict, List, Literal, Mapping, Tuple
+from beartype.typing import Any, Dict, List, Literal, Mapping, Tuple
 
 import numpy as np
 import polars as pl
diff --git a/alphafold3_pytorch/data/weighted_pdb_sampler.py b/alphafold3_pytorch/data/weighted_pdb_sampler.py
index 4ae749dd..74fdce5d 100644
--- a/alphafold3_pytorch/data/weighted_pdb_sampler.py
+++ b/alphafold3_pytorch/data/weighted_pdb_sampler.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import os
-from typing import Dict, Iterator, List, Literal, Tuple
+from beartype.typing import Dict, Iterator, List, Literal, Tuple
 
 import numpy as np
 import polars as pl
diff --git a/alphafold3_pytorch/inputs.py b/alphafold3_pytorch/inputs.py
index d6441295..afbc7fce 100644
--- a/alphafold3_pytorch/inputs.py
+++ b/alphafold3_pytorch/inputs.py
@@ -14,7 +14,16 @@
 from itertools import groupby
 from pathlib import Path
 from retrying import retry
-from typing import Any, Callable, Dict, List, Literal, Set, Tuple, Type
+from beartype.typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Set,
+    Tuple,
+    Type,
+)
 
 import einx
 import numpy as np
diff --git a/alphafold3_pytorch/life.py b/alphafold3_pytorch/life.py
index e37b02e9..cdb5c808 100644
--- a/alphafold3_pytorch/life.py
+++ b/alphafold3_pytorch/life.py
@@ -1,5 +1,5 @@
 import os
-from typing import Literal
+from beartype.typing import Literal
 
 import gemmi
 import rdkit.Geometry.rdGeometry as rdGeometry
diff --git a/alphafold3_pytorch/trainer.py b/alphafold3_pytorch/trainer.py
index 8e77a937..3e118c83 100644
--- a/alphafold3_pytorch/trainer.py
+++ b/alphafold3_pytorch/trainer.py
@@ -7,7 +7,7 @@
 from alphafold3_pytorch.alphafold3 import Alphafold3
 from alphafold3_pytorch.attention import pad_at_dim, pad_or_slice_to
 
-from typing import TypedDict, List, Callable
+from beartype.typing import TypedDict, List, Callable
 
 from alphafold3_pytorch.tensor_typing import (
     should_typecheck,
diff --git a/alphafold3_pytorch/utils/data_utils.py b/alphafold3_pytorch/utils/data_utils.py
index e1084e2b..d0410813 100644
--- a/alphafold3_pytorch/utils/data_utils.py
+++ b/alphafold3_pytorch/utils/data_utils.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Literal, Set, Tuple
+from beartype.typing import Any, Dict, List, Literal, Set, Tuple
 
 import numpy as np
 import torch
diff --git a/alphafold3_pytorch/utils/model_utils.py b/alphafold3_pytorch/utils/model_utils.py
index 53468acd..9daebb27 100644
--- a/alphafold3_pytorch/utils/model_utils.py
+++ b/alphafold3_pytorch/utils/model_utils.py
@@ -1,8 +1,8 @@
 from functools import wraps
-from typing import Callable, List, Tuple, Union
+from beartype.typing import Callable, List, Tuple, Union
 
 import einx
-import pkg_resources
+import importlib.metadata
 import torch
 import torch.nn.functional as F
 from einops import einsum, pack, rearrange, reduce, repeat, unpack
@@ -644,8 +644,9 @@ def package_available(package_name: str) -> bool:
     :return: `True` if the package is available. `False` otherwise.
     """
     try:
-        return pkg_resources.require(package_name) is not None
-    except pkg_resources.DistributionNotFound:
+        importlib.metadata.version(package_name)
+        return True
+    except importlib.metadata.PackageNotFoundError:
         return False
 
 
diff --git a/alphafold3_pytorch/utils/utils.py b/alphafold3_pytorch/utils/utils.py
index 6c4e17ee..4172a908 100644
--- a/alphafold3_pytorch/utils/utils.py
+++ b/alphafold3_pytorch/utils/utils.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from typing import Any, Iterable, List
+from beartype.typing import Any, Iterable, List
 
 
 def exists(val: Any) -> bool:
diff --git a/pyproject.toml b/pyproject.toml
index 3c58848b..70dc1683 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -43,7 +43,7 @@ dependencies = [
     "huggingface_hub>=0.21.4",
     "jaxtyping>=0.2.28",
     "lightning>=2.2.5",
-    "numpy==1.23.5",
+    "numpy>=1.23.5",
     "polars>=1.1.0",
     "pdbeccdutils>=0.8.5",
     "pydantic>=2.8.2",
diff --git a/scripts/cluster_pdb_test_mmcifs.py b/scripts/cluster_pdb_test_mmcifs.py
index 6f1168dd..8d0b6d2b 100644
--- a/scripts/cluster_pdb_test_mmcifs.py
+++ b/scripts/cluster_pdb_test_mmcifs.py
@@ -24,7 +24,7 @@
 import os
 from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor, as_completed
-from typing import Dict, List, Set, Tuple
+from beartype.typing import Dict, List, Set, Tuple
 
 import numpy as np
 import polars as pl
diff --git a/scripts/cluster_pdb_train_mmcifs.py b/scripts/cluster_pdb_train_mmcifs.py
index f0004f4f..b7a40cad 100644
--- a/scripts/cluster_pdb_train_mmcifs.py
+++ b/scripts/cluster_pdb_train_mmcifs.py
@@ -25,7 +25,7 @@
 import os
 import subprocess
 from concurrent.futures import ProcessPoolExecutor, as_completed
-from typing import Dict, List, Literal, Optional, Set, Tuple, Union
+from beartype.typing import Dict, List, Literal, Optional, Set, Tuple, Union
 
 import numpy as np
 import polars as pl
diff --git a/scripts/cluster_pdb_val_mmcifs.py b/scripts/cluster_pdb_val_mmcifs.py
index b5cd8e38..5b89668e 100644
--- a/scripts/cluster_pdb_val_mmcifs.py
+++ b/scripts/cluster_pdb_val_mmcifs.py
@@ -38,7 +38,7 @@
 import subprocess
 from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor, as_completed
-from typing import Dict, List, Set, Tuple
+from beartype.typing import Dict, List, Set, Tuple
 
 import numpy as np
 import polars as pl