build(python): Do not set environment variable on import

crates/polars-core/src/chunked_array/object/extension/mod.rs

@@ -3,6 +3,7 @@ mod list;
 pub(crate) mod polars_extension;
 
 use std::mem;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 use arrow::array::FixedSizeBinaryArray;
 use arrow::bitmap::MutableBitmap;
@@ -13,6 +14,14 @@ use crate::prelude::*;
 use crate::PROCESS_ID;
 
 pub const EXTENSION_NAME: &str = "POLARS_EXTENSION_TYPE";
+static POLARS_ALLOW_EXTENSION: AtomicBool = AtomicBool::new(false);
+
+/// Control whether extension types may be created.
+///
+/// If the environment variable POLARS_ALLOW_EXTENSION is set, this function has no effect.
+pub fn set_polars_allow_extension(toggle: bool) {
+    POLARS_ALLOW_EXTENSION.store(toggle, Ordering::Relaxed)
+}
 
 /// Invariants
 /// `ptr` must point to start a `T` allocation
@@ -54,9 +63,9 @@ pub(crate) fn create_extension<I: Iterator<Item = Option<T>> + TrustedLen, T: Si
     iter: I,
 ) -> PolarsExtension {
     let env = "POLARS_ALLOW_EXTENSION";
-    std::env::var(env).unwrap_or_else(|_| {
-        panic!("env var: {env} must be set to allow extension types to be created",)
-    });
+    if !(POLARS_ALLOW_EXTENSION.load(Ordering::Relaxed) || std::env::var(env).is_ok()) {
+        panic!("creating extension types not allowed - try setting the environment variable {env}")
+    }
     let t_size = std::mem::size_of::<T>();
     let t_alignment = std::mem::align_of::<T>();
     let n_t_vals = iter.size_hint().1.unwrap();
@@ -176,7 +185,7 @@ mod test {
 
     #[test]
     fn test_create_extension() {
-        std::env::set_var("POLARS_ALLOW_EXTENSION", "1");
+        set_polars_allow_extension(true);
         // Run this under MIRI.
         let foo = Foo {
             a: 1,
@@ -195,7 +204,7 @@ mod test {
 
     #[test]
     fn test_extension_to_list() {
-        std::env::set_var("POLARS_ALLOW_EXTENSION", "1");
+        set_polars_allow_extension(true);
         let foo1 = Foo {
             a: 1,
             b: 1,
@@ -219,7 +228,7 @@ mod test {
 
     #[test]
     fn test_extension_to_list_explode() {
-        std::env::set_var("POLARS_ALLOW_EXTENSION", "1");
+        set_polars_allow_extension(true);
         let foo1 = Foo {
             a: 1,
             b: 1,

crates/polars-core/src/chunked_array/object/mod.rs

@@ -15,6 +15,8 @@ mod is_valid;
 mod iterator;
 pub mod registry;
 
+pub use extension::set_polars_allow_extension;
+
 #[derive(Debug, Clone)]
 pub struct ObjectArray<T>
 where

py-polars/polars/__init__.py

@@ -1,5 +1,4 @@
 import contextlib
-import os
 
 with contextlib.suppress(ImportError):  # Module not available when building docs
     # ensure the object constructor is known by polars
@@ -13,9 +12,7 @@
     # we also set other function pointers needed
     # on the rust side. This function is highly
     # unsafe and should only be called once.
-    from polars.polars import (
-        __register_startup_deps,
-    )
+    from polars.polars import __register_startup_deps
 
     __register_startup_deps()
 
@@ -381,8 +378,6 @@
     "sql_expr",
 ]
 
-os.environ["POLARS_ALLOW_EXTENSION"] = "true"
-
 
 def __getattr__(name: str):  # type: ignore[no-untyped-def]
     # Deprecate re-export of exceptions at top-level

py-polars/src/on_startup.rs

@@ -2,8 +2,8 @@ use std::any::Any;
 
 use polars::prelude::*;
 use polars_core::chunked_array::object::builder::ObjectChunkedBuilder;
-use polars_core::chunked_array::object::registry;
 use polars_core::chunked_array::object::registry::AnonymousObjectBuilder;
+use polars_core::chunked_array::object::{registry, set_polars_allow_extension};
 use polars_core::error::PolarsError::ComputeError;
 use polars_error::PolarsWarning;
 use pyo3::intern;
@@ -67,6 +67,7 @@ fn warning_function(msg: &str, warning: PolarsWarning) {
 
 #[pyfunction]
 pub fn __register_startup_deps() {
+    set_polars_allow_extension(true);
     if !registry::is_object_builder_registered() {
         // Stack frames can get really large in debug mode.
         #[cfg(debug_assertions)]