From f87c7819fb0c0865bf4b3314525b473d550a4e27 Mon Sep 17 00:00:00 2001
From: Ben Dowling <ben.dowling@icloud.com>
Date: Tue, 26 Apr 2022 22:32:06 +0100
Subject: [PATCH] Add itertools to modules (#5140)

* GH hygiene: contributing guide, templates, stalebot (#4967)

* Update contributing guide

* Update issue + PR templates

* Stalebot for all issues, no exceptions

* Update links

* Missed one

* PR feedback

* Update CHANGELOG

Co-authored-by: Jeremy Cohen <jeremy@dbtlabs.com>
---
 .../unreleased/Features-20220424-132655.yaml  |  7 ++++++
 core/dbt/context/base.py                      | 25 +++++++++++++++++++
 2 files changed, 32 insertions(+)
 create mode 100644 .changes/unreleased/Features-20220424-132655.yaml

diff --git a/.changes/unreleased/Features-20220424-132655.yaml b/.changes/unreleased/Features-20220424-132655.yaml
new file mode 100644
index 00000000000..04984324e60
--- /dev/null
+++ b/.changes/unreleased/Features-20220424-132655.yaml
@@ -0,0 +1,7 @@
+kind: Features
+body: Adds itertools to modules Jinja namespace
+time: 2022-04-24T13:26:55.008246+01:00
+custom:
+  Author: bd3dowling
+  Issue: "5130"
+  PR: "5140"
diff --git a/core/dbt/context/base.py b/core/dbt/context/base.py
index 355937bb4fd..c2bb6417304 100644
--- a/core/dbt/context/base.py
+++ b/core/dbt/context/base.py
@@ -23,6 +23,7 @@
 import pytz
 import datetime
 import re
+import itertools
 
 # Contexts in dbt Core
 # Contexts are used for Jinja rendering. They include context methods,
@@ -77,11 +78,35 @@ def get_re_module_context() -> Dict[str, Any]:
     return {name: getattr(re, name) for name in context_exports}
 
 
+def get_itertools_module_context() -> Dict[str, Any]:
+    # Excluded dropwhile, filterfalse, takewhile and groupby;
+    # first 3 illogical for Jinja and last redundant.
+    context_exports = [
+        "count",
+        "cycle",
+        "repeat",
+        "accumulate",
+        "chain",
+        "compress",
+        "islice",
+        "starmap",
+        "tee",
+        "zip_longest",
+        "product",
+        "permutations",
+        "combinations",
+        "combinations_with_replacement",
+    ]
+
+    return {name: getattr(itertools, name) for name in context_exports}
+
+
 def get_context_modules() -> Dict[str, Dict[str, Any]]:
     return {
         "pytz": get_pytz_module_context(),
         "datetime": get_datetime_module_context(),
         "re": get_re_module_context(),
+        "itertools": get_itertools_module_context(),
     }