Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use update instead of assignment to set _charges #1069

Merged
merged 1 commit into from
Oct 2, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions openff/interchange/common/_nonbonded.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import abc
from collections.abc import Iterable
from typing import Literal
from typing import Any, Literal

from openff.toolkit import Quantity, unit
from pydantic import Field, PrivateAttr, computed_field
Expand Down Expand Up @@ -101,8 +101,7 @@ class ElectrostaticsCollection(_NonbondedCollection):
nonperiodic_potential: Literal["Coulomb", "cutoff", "no-cutoff"] = Field("Coulomb")
exception_potential: Literal["Coulomb"] = Field("Coulomb")

# TODO: Charge caching doesn't work when this is defined in the model
# _charges: dict[Any, _ElementaryChargeQuantity] = PrivateAttr(default_factory=dict)
_charges: dict[Any, _ElementaryChargeQuantity] = PrivateAttr(default_factory=dict)
_charges_cached: bool = PrivateAttr(default=False)

@computed_field # type: ignore[misc]
Expand All @@ -112,7 +111,7 @@ def charges(
) -> dict[TopologyKey | LibraryChargeTopologyKey | VirtualSiteKey, _ElementaryChargeQuantity]:
"""Get the total partial charge on each atom, including virtual sites."""
if len(self._charges) == 0 or self._charges_cached is False: # type: ignore[has-type]
self._charges = self._get_charges(include_virtual_sites=False)
self._charges.update(self._get_charges(include_virtual_sites=False))
Comment on lines -115 to +114
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This ... makes a fair amount of sense to me, actually. I think in non-Pydantic world it's best practice to update a dict instead of re-assign it, but I'm also pretty sure something in Pydantic's magic machinery means just passing the return value of another function to an instance attribute gummed something up something.

self._charges_cached = True

return self._charges
Expand Down
5 changes: 2 additions & 3 deletions openff/interchange/smirnoff/_nonbonded.py
Original file line number Diff line number Diff line change
Expand Up @@ -273,8 +273,7 @@ class SMIRNOFFElectrostaticsCollection(ElectrostaticsCollection, SMIRNOFFCollect
) # type: ignore[assignment]
exception_potential: Literal["Coulomb"] = Field("Coulomb")

# TODO: Charge caching doesn't work when this is defined in the model
# _charges: dict[Any, _ElementaryChargeQuantity] = PrivateAttr(default_factory=dict)
_charges: dict[Any, _ElementaryChargeQuantity] = PrivateAttr(default_factory=dict)
_charges_cached: bool = PrivateAttr(default=False)

@classmethod
Expand Down Expand Up @@ -305,7 +304,7 @@ def charges(
) -> dict[TopologyKey | LibraryChargeTopologyKey | VirtualSiteKey, _ElementaryChargeQuantity]:
"""Get the total partial charge on each atom, including virtual sites."""
if len(self._charges) == 0 or self._charges_cached is False:
self._charges = self._get_charges(include_virtual_sites=True)
self._charges.update(self._get_charges(include_virtual_sites=True))
self._charges_cached = True

return self._charges
Expand Down